#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V7 — Column-aware Y-bucket).
Thuật toán:
  1. Phát hiện layout nhiều cột qua X-histogram (gap detection)
  2. Nếu 2+ cột thực sự (balanced) → đọc column-by-column
  3. Nếu 1 cột → Y-bucket row-by-row (int(y/threshold), sort by X)
  4. Split multi-line blocks → từng dòng riêng

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

import os
import sys
import json
import glob
import re
from collections import defaultdict

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

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

# Y-threshold for row grouping: items within this many Y-units (normalized)
# are considered on the same line. ~0.022 ≈ 1 line of Myanmar text at 12pt
Y_THRESHOLD = 0.025

# Column detection
MIN_COL_GAP = 0.08         # Min X gap to split columns
MIN_COL_FRAGS = 4          # Min fragments per column for real multi-column
COL_BALANCE_RATIO = 0.25   # Min ratio of smaller/larger column to be "balanced"

# X gap spacing in output
COL_GAP = 0.15
WORD_GAP = 0.03


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

def extract_text(obj):
    """Extract text 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):
    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."""
    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:
            line_h = (y_max - y_min) / len(lines)
            for i, line in enumerate(lines):
                result.append((x_min, y_min + i*line_h, x_max, 
                              y_min + (i+1)*line_h, line))
    return result


def detect_columns(fragments):
    """Detect if page has balanced multi-column layout.
    
    Uses X-midpoint histogram: find gaps, check if columns are balanced.
    Returns list of column X-ranges [(x_min, x_max), ...] or [(0,1)].
    """
    if len(fragments) < 12:
        return [(0, 1)]
    
    # Get X-mids of narrow fragments only (skip full-width)
    x_data = []
    for fx_min, fy_min, fx_max, fy_max, text in fragments:
        width = fx_max - fx_min
        if width > 0.60:
            continue
        x_data.append((fx_min + fx_max) / 2)
    
    if len(x_data) < 10:
        return [(0, 1)]
    
    x_data.sort()
    
    # Find gaps
    gaps = []
    for i in range(len(x_data) - 1):
        g = x_data[i+1] - x_data[i]
        if g > MIN_COL_GAP:
            gaps.append((i, g, x_data[i], x_data[i+1]))
    
    if not gaps:
        return [(0, 1)]
    
    # Score each gap: prefer balanced splits
    best = None
    best_score = 0
    for idx, gap, left_x, right_x in gaps:
        left_n = idx + 1
        right_n = len(x_data) - idx - 1
        
        if left_n >= MIN_COL_FRAGS and right_n >= MIN_COL_FRAGS:
            balance = min(left_n, right_n) / max(left_n, right_n)
            if balance >= COL_BALANCE_RATIO:
                score = gap * 10 + balance * 5
                if score > best_score:
                    best_score = score
                    best = (left_x, right_x)
    
    if best is None:
        return [(0, 1)]
    
    boundary = (best[0] + best[1]) / 2
    
    # Compute column ranges
    col1_frags = [f for f in fragments if (f[0] + f[2])/2 < boundary and (f[2]-f[0]) < 0.60]
    col2_frags = [f for f in fragments if (f[0] + f[2])/2 >= boundary and (f[2]-f[0]) < 0.60]
    
    if len(col1_frags) < MIN_COL_FRAGS or len(col2_frags) < MIN_COL_FRAGS:
        return [(0, 1)]
    
    col1_xmin = min(f[0] for f in col1_frags)
    col1_xmax = max(f[2] for f in col1_frags)
    col2_xmin = min(f[0] for f in col2_frags)
    col2_xmax = max(f[2] for f in col2_frags)
    
    return [(col1_xmin, col1_xmax), (col2_xmin, col2_xmax)]


def y_bucket_key(y):
    """Map Y to bucket using int(y/threshold) as suggested."""
    return int(y / Y_THRESHOLD)


def render_bucket(fragments):
    """Render fragments in a bucket as one line, sorted left-to-right."""
    if not fragments:
        return ''
    sorted_frags = sorted(fragments, key=lambda f: f[0])
    parts = []
    prev_xmax = 0
    for i, (xmin, ymin, xmax, ymax, text) in enumerate(sorted_frags):
        if i > 0:
            gap = xmin - prev_xmax
            if gap > COL_GAP:
                parts.append('    ')
            elif gap > WORD_GAP:
                parts.append('  ')
            else:
                parts.append(' ')
        parts.append(text)
        prev_xmax = xmax
    return ''.join(parts)


def render_column(fragments):
    """Render fragments within one column: Y-bucket → lines."""
    buckets = defaultdict(list)
    for f in fragments:
        ymid = (f[1] + f[3]) / 2
        buckets[y_bucket_key(ymid)].append(f)
    
    lines = []
    for bk in sorted(buckets.keys()):
        line = render_bucket(buckets[bk])
        if line:
            lines.append(line)
    return lines


def page_to_text(page_data, page_num, debug=False):
    blocks = page_data.get('blocks', [])
    
    # Extract paragraphs
    all_paras = []
    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        for para in block.get('paragraphs', []):
            text = extract_text(para).strip()
            if not text:
                continue
            x_min, y_min, x_max, y_max = get_bbox(para)
            if is_scanner(text):
                continue
            all_paras.append((x_min, y_min, x_max, y_max, text))
    
    if not all_paras:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n"
    
    # Split multi-line
    fragments = split_into_lines(all_paras)
    
    # Separate full-width from narrow fragments
    full_width = [f for f in fragments if (f[2] - f[0]) > 0.60]
    narrow = [f for f in fragments if (f[2] - f[0]) <= 0.60]
    
    # Detect columns on narrow fragments
    columns = detect_columns(narrow)
    
    output_lines = [f"## PAGE {page_num}"]
    if debug:
        output_lines.append(f"<!-- {len(fragments)} frags, {len(columns)} col(s), {len(full_width)} full-width -->")
    output_lines.append("")
    
    if len(columns) == 1:
        # Single column: render all fragments together
        lines = render_column(fragments)
        output_lines.extend(lines)
    else:
        # Multi-column: render each column separately
        # Full-width fragments act as section breaks at their Y position
        used_frags = set()
        
        for col_idx, (col_xmin, col_xmax) in enumerate(columns):
            col_frags = []
            for f in narrow:
                xmid = (f[0] + f[2]) / 2
                if col_xmin - 0.02 <= xmid <= col_xmax + 0.02:
                    col_frags.append(f)
                    used_frags.add(id(f))
            
            if col_idx > 0:
                output_lines.append("")
            
            col_lines = render_column(col_frags)
            output_lines.extend(col_lines)
        
        # Handle any leftover fragments
        leftover = [f for f in narrow if id(f) not in used_frags]
        if leftover:
            output_lines.append("")
            output_lines.extend(render_column(leftover))
    
    # Full-width fragments at their natural Y positions
    if full_width:
        output_lines.append("")
        output_lines.extend(render_column(full_width))
    
    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_v7.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(30))
            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()