#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V6 — Y-bucket grouping).
- Split multi-line blocks → các dòng riêng
- Group fragment theo Y-bucket (fixed bin) → tránh overlap chain
- Sort X trong mỗi bucket → trái→phải
- Không cần column detection phức tạp — Y-bucket tự xử lý mọi layout

Usage:
  python3 json_to_markdown_v6.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),
]

# Y-bucket size (normalized, ~typical line height for Myanmar text)
BUCKET_SIZE = 0.022

# X gap threshold for column separator (4 spaces vs 2 spaces)
COLUMN_GAP_THRESHOLD = 0.15
WORD_GAP_THRESHOLD = 0.03


# ─── Core ────────────────────────────────────────────────────────────────────

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 get_bbox(obj):
    """Get (x_min, y_min, x_max, y_max) from normalizedVertices."""
    verts = obj.get('boundingBox', {}).get('normalizedVertices', [])
    if not verts:
        return (0, 0, 0, 0)
    return (
        min(v.get('x', 0) for v in verts),
        min(v.get('y', 0) for v in verts),
        max(v.get('x', 0) for v in verts),
        max(v.get('y', 0) for v in verts),
    )


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


def split_into_lines(paragraphs):
    """Split multi-line paragraphs into individual line fragments.
    Returns list of (x_min, y_min, x_max, y_max, text_line).
    """
    result = []
    for x_min, y_min, x_max, y_max, text in paragraphs:
        lines = text.split('\n')
        lines = [l.strip() for l in lines if l.strip()]
        if not lines:
            continue
        
        if len(lines) == 1:
            result.append((x_min, y_min, x_max, y_max, lines[0]))
        else:
            num_lines = len(lines)
            line_h = (y_max - y_min) / num_lines
            for i, line in enumerate(lines):
                ly_min = y_min + i * line_h
                ly_max = ly_min + line_h
                result.append((x_min, ly_min, x_max, ly_max, line))
    return result


def y_bucket_key(y_min, y_max):
    """Map a Y range to a bucket key based on its midpoint."""
    y_mid = (y_min + y_max) / 2
    return round(y_mid / BUCKET_SIZE) * BUCKET_SIZE


def bucket_to_text(fragments):
    """Convert fragments in a single Y-bucket to a text line.
    Fragments are sorted by X (left to right) with column/word spacing.
    """
    if not fragments:
        return ''
    
    sorted_frags = sorted(fragments, key=lambda f: f[0])
    parts = []
    prev_x_max = 0
    
    for i, (x_min, y_min, x_max, y_max, text) in enumerate(sorted_frags):
        if i > 0:
            gap = x_min - prev_x_max
            if gap > COLUMN_GAP_THRESHOLD:
                parts.append('    ')
            elif gap > WORD_GAP_THRESHOLD:
                parts.append('  ')
            else:
                parts.append(' ')
        parts.append(text)
        prev_x_max = x_max
    
    return ''.join(parts)


def page_to_text(page_data, page_num, debug=False):
    """Convert a single Cloud Vision page to Markdown text."""
    blocks = page_data.get('blocks', [])
    
    # Step 1: Extract all paragraphs
    all_paragraphs = []
    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
            
            x_min, y_min, x_max, y_max = get_bbox(para)
            
            if is_scanner(text):
                continue
            
            all_paragraphs.append((x_min, y_min, x_max, y_max, text))
    
    if not all_paragraphs:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n"
    
    # Step 2: Split multi-line paragraphs into individual lines
    fragments = split_into_lines(all_paragraphs)
    
    # Step 3: Group fragments into Y-buckets
    buckets = {}
    for frag in fragments:
        key = y_bucket_key(frag[1], frag[3])
        if key not in buckets:
            buckets[key] = []
        buckets[key].append(frag)
    
    # Step 4: Build output — one line per bucket, sorted by Y
    output_lines = [f"## PAGE {page_num}"]
    
    if debug:
        output_lines.append(f"<!-- {len(fragments)} fragments, {len(buckets)} buckets -->")
    
    output_lines.append("")
    
    for bk in sorted(buckets.keys()):
        line = bucket_to_text(buckets[bk])
        if line:
            output_lines.append(line)
    
    output_lines.append("")
    return '\n'.join(output_lines)


def json_to_markdown(json_path, output_dir, start_page, debug=False):
    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 = []
    for i, response in enumerate(data.get('responses', [])):
        page_num = start_page + i
        fta = response.get('fullTextAnnotation', {})
        pages = fta.get('pages', [])
        
        if pages:
            page_text = page_to_text(pages[0], page_num, debug=debug)
        else:
            raw_text = fta.get('text', '').strip()
            page_text = f"## PAGE {page_num}\n\n{raw_text if raw_text else '*(trang trống)*'}\n"
        page_texts.append(page_text)

    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)


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v6.py <json_dir> <output_dir> [--dry-run] [--debug]")
        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"📄 Tìm thấy {len(json_files)} file JSON")
    total_pages = 0

    for jf in json_files:
        out_path, pages = json_to_markdown(jf, output_dir, total_pages + 1, debug=debug)
        total_pages += pages
        
        if dry_run:
            with open(out_path, 'r') as f:
                preview = ''.join(f.readline() for _ in range(25))
            print(f"\n{'='*60}")
            print(f"📄 {os.path.basename(out_path)} ({pages} trang) — PREVIEW:")
            print(f"{'='*60}")
            print(preview)
            print(f"... (saved at {out_path})")
        else:
            print(f"   ✅ {os.path.basename(out_path)} ({pages} trang)")

    if not dry_run:
        print(f"\n🎉 Done! {len(json_files)} file .md → {os.path.abspath(output_dir)}")
        print(f"   Tổng: {total_pages} trang")


if __name__ == '__main__':
    main()