#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V9 — Raw Paragraph Extraction)

KHÔNG clean header/footer — chỉ extract nguyên vẹn text theo thứ tự paragraph từ Cloud Vision.
Clean header/artifacts sẽ được làm ở script riêng, chạy sau khi đã có file .md.

Khác biệt so với V7/V8:
- V7: fragment-level Y-bucket + header/footer cleanup → layout đẹp nhưng có thể xáo trộn text
- V8: fragment-level Y-bucket, chỉ remove URL → vẫn xáo trộn text ở vùng ranh giới header
- V9: paragraph-level extraction, KHÔNG filter → giữ nguyên thứ tự & nội dung từ Cloud Vision

Usage:
  python3 json_to_markdown_v9.py <json_dir> <output_dir> [--dry-run] [--debug]
"""

import os
import sys
import json
import glob
import re

# ─── Constants ───────────────────────────────────────────────────────────────

SCANNER_PATTERNS = [
    re.compile(r'scanned\s+with', re.IGNORECASE),
    re.compile(r'cam\s*scanner', re.IGNORECASE),
]


# ─── Text Extraction ─────────────────────────────────────────────────────────

def extract_text(obj):
    """Extract text from a paragraph/word, preserving spaces and line breaks."""
    text = ''
    for word in obj.get('words', []):
        for sym in word.get('symbols', []):
            text += sym.get('text', '')
            prop = sym.get('property', {})
            if prop.get('detectedBreak'):
                bt = prop['detectedBreak']['type']
                if bt in ('SPACE', 'SURE_SPACE'):
                    text += ' '
                elif bt in ('EOL_SURE_SPACE', 'LINE_BREAK'):
                    text += '\n'
    return text


def is_scanner(text):
    for pat in SCANNER_PATTERNS:
        if pat.search(text):
            return True
    return False


# ─── Page Processing ─────────────────────────────────────────────────────────

def page_to_text(page_data, page_num, debug=False):
    """Convert a single Cloud Vision page to raw Markdown text.
    
    Preserves original paragraph order from Cloud Vision — no Y-bucket,
    no header/footer filtering, no fragment rearrangement.
    """
    blocks = page_data.get('blocks', [])
    
    paragraphs = []
    para_count = 0
    
    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        for para in block.get('paragraphs', []):
            text = extract_text(para)
            text = text.strip()
            if not text:
                continue
            if is_scanner(text):
                continue
            paragraphs.append(text)
            para_count += 1
    
    if not paragraphs:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n"
    
    # Build output: page header + paragraphs separated by blank lines
    output_lines = [f"## PAGE {page_num}"]
    if debug:
        output_lines.append(f"<!-- {para_count} paragraphs, raw extraction -->")
    output_lines.append("")
    
    for para_text in paragraphs:
        output_lines.append(para_text)
        output_lines.append("")
    
    return '\n'.join(output_lines)


def json_to_markdown(json_path, output_dir, start_page, debug=False):
    """Process one JSON file → raw Markdown file."""
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    basename = os.path.splitext(os.path.basename(json_path))[0]
    out_path = os.path.join(output_dir, f"{basename}.md")

    page_texts = []
    para_stats = []
    
    for i, response in enumerate(data.get('responses', [])):
        page_num = start_page + i
        fta = response.get('fullTextAnnotation', {})

        if 'pages' not in fta:
            raw_text = fta.get('text', '').strip()
            page_texts.append(
                f"## PAGE {page_num}\n\n{raw_text if raw_text else '*(trang trống)*'}\n"
            )
            para_stats.append(0)
            continue

        pages = fta.get('pages', [])
        if pages:
            page_text = page_to_text(pages[0], page_num, debug=debug)
            page_texts.append(page_text)
            # Count paragraphs
            para_count = page_text.count('\n\n') - 1  # subtract page header
            para_stats.append(max(0, para_count))
            if debug:
                print(f"   Page {page_num:2d}  {para_count} paragraphs")
        else:
            page_texts.append(f"## PAGE {page_num}\n\n*(trang trống)*\n")
            para_stats.append(0)

    os.makedirs(output_dir, exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(page_texts))

    return out_path, len(page_texts), sum(para_stats)


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v9.py <json_dir> <output_dir> [--dry-run] [--debug]")
        print()
        print("  json_dir    — thư mục chứa các file JSON từ Cloud Vision")
        print("  output_dir  — thư mục xuất file Markdown")
        print("  --dry-run   — preview 25 dòng đầu mỗi file, không lưu thật")
        print("  --debug     — hiển thị paragraph count mỗi trang")
        sys.exit(1)

    json_dir = sys.argv[1]
    output_dir = sys.argv[2]
    dry_run = '--dry-run' in sys.argv
    debug = '--debug' in sys.argv

    json_files = sorted(
        glob.glob(os.path.join(json_dir, '*.json')),
        key=lambda f: int(re.search(r'output-(\d+)-to', os.path.basename(f)).group(1))
    )

    if not json_files:
        print(f"❌ Không tìm thấy file JSON nào trong {json_dir}")
        sys.exit(1)

    print(f"🔧 V9: Raw paragraph extraction (no cleanup, no Y-bucket)")
    print(f"📄 {len(json_files)} file JSON tìm thấy\n")

    total_pages = total_paras = 0
    all_md = []

    for jf in json_files:
        out_path, pages, paras = json_to_markdown(jf, output_dir, total_pages + 1, debug=debug)
        total_pages += pages
        total_paras += paras

        if dry_run:
            with open(out_path, 'r') as f:
                preview = ''.join(f.readline() for _ in range(30))
            print(f"\n{'='*60}")
            print(f"📄 {os.path.basename(out_path)} ({pages} trang, {paras} paragraphs) — PREVIEW:")
            print(f"{'='*60}")
            print(preview)
            print(f"... (saved at {out_path})")
        else:
            print(f"   ✅ {os.path.basename(out_path)} ({pages} trang, {paras} paragraphs)")

        with open(out_path, 'r') as f:
            all_md.append(f.read())

    if not dry_run:
        merged_path = os.path.join(output_dir, "full-merged.md")
        with open(merged_path, 'w', encoding='utf-8') as f:
            f.write(f"# OCR Raw Extraction (V9: Paragraph-level, no cleanup)\n\n")
            f.write('\n'.join(all_md))
            f.write(f"\n\n---\n**Tổng:** {total_pages} trang | {total_paras} paragraphs\n")
        print(f"\n📄 Merged: {merged_path} ({os.path.getsize(merged_path):,} bytes)")

    print(f"\n{'='*60}")
    print(f"📊 {total_pages} trang | {total_paras} paragraphs (raw — chưa clean header/artifacts)")
    print(f"🎉 Done! → {os.path.abspath(output_dir)}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()