#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V2 — Sắp xếp theo tọa độ).
Khắc phục vấn đề layout hỗn loạn: sắp xếp paragraph theo Y (trên→dưới),
sau đó X (trái→phải), group theo hàng dựa trên Y overlap.

Usage:
  python3 json_to_markdown_v2.py <json_dir> <output_dir> [--dry-run]
Example:
  python3 json_to_markdown_v2.py "010-pali-thaykha/ocr/raw/" "010-pali-thaykha/extracted/"
  python3 json_to_markdown_v2.py "010-pali-thaykha/ocr/raw/" "010-pali-thaykha/extracted/" --dry-run
"""

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

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

# Scanner watermark patterns (case-insensitive)
SCANNER_PATTERNS = [
    re.compile(r'scanned\s+with', re.IGNORECASE),
    re.compile(r'cam\s*scanner', re.IGNORECASE),
    re.compile(r'CS\s*CamScanner', re.IGNORECASE),
]

# Myanmar header patterns (page headers/footers that are not content)
HEADER_PATTERNS_MY = [
    re.compile(r'^စာမျက်နှာ\s*$'),   # "page" label
    re.compile(r'^မာတိကာခေါင်းစဉ်\s*$'),  # "table of contents heading" label
    re.compile(r'^အကြောင်းအရာ\s*$'),  # "content" label
]

# Maximum Y for header zone (normalized 0-1)
HEADER_Y_MAX = 0.12
# Minimum Y for footer zone
FOOTER_Y_MIN = 0.88


# ─── Core Functions ──────────────────────────────────────────────────────────

def extract_paragraph_text(para):
    """Extract text from a paragraph, preserving spaces and line breaks."""
    text = ''
    for word in para.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_bounding_box(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)
    x_min = min(v.get('x', 0) for v in verts)
    y_min = min(v.get('y', 0) for v in verts)
    x_max = max(v.get('x', 0) for v in verts)
    y_max = max(v.get('y', 0) for v in verts)
    return (x_min, y_min, x_max, y_max)


def is_scanner_artifact(text):
    """Check if text is a scanner watermark."""
    text_stripped = text.strip()
    for pat in SCANNER_PATTERNS:
        if pat.search(text_stripped):
            return True
    return False


def is_header_label(text, y_max):
    """Check if text is a header label (not real content)."""
    text_stripped = text.strip()
    # Check if in header zone
    if y_max < HEADER_Y_MAX:
        for pat in HEADER_PATTERNS_MY:
            if pat.match(text_stripped):
                return True
    return False


def is_footer_artifact(y_min):
    """Check if paragraph is in footer zone."""
    return y_min > FOOTER_Y_MIN


def y_overlap(y1_min, y1_max, y2_min, y2_max, threshold=0.3):
    """Check if two Y ranges overlap significantly.
    
    Returns fraction of overlap relative to the smaller range.
    """
    overlap_start = max(y1_min, y2_min)
    overlap_end = min(y1_max, y2_max)
    overlap = max(0, overlap_end - overlap_start)
    
    range1 = y1_max - y1_min
    range2 = y2_max - y2_min
    min_range = min(range1, range2)
    
    if min_range == 0:
        return 0
    return overlap / min_range


def group_paragraphs_into_rows(paragraphs):
    """Group paragraphs into rows based on Y overlap.
    
    Each paragraph is (x_min, y_min, x_max, y_max, text).
    Returns list of rows, each row is a list of paragraphs sorted by X.
    """
    if not paragraphs:
        return []
    
    # Sort by y_min (top to bottom)
    sorted_paras = sorted(paragraphs, key=lambda p: (p[1], p[0]))
    
    rows = []
    current_row = [sorted_paras[0]]
    current_y_min = sorted_paras[0][1]
    current_y_max = sorted_paras[0][3]
    
    for para in sorted_paras[1:]:
        _, y_min, _, y_max, _ = para
        
        # Check if this paragraph overlaps with current row's Y range
        overlap = y_overlap(current_y_min, current_y_max, y_min, y_max, threshold=0.2)
        
        if overlap > 0:
            # Same row: expand Y range
            current_row.append(para)
            current_y_min = min(current_y_min, y_min)
            current_y_max = max(current_y_max, y_max)
        else:
            # New row
            # Sort current row by X (left to right)
            current_row.sort(key=lambda p: p[0])
            rows.append(current_row)
            current_row = [para]
            current_y_min = y_min
            current_y_max = y_max
    
    # Don't forget the last row
    if current_row:
        current_row.sort(key=lambda p: p[0])
        rows.append(current_row)
    
    return rows


def page_to_text(page_data, page_num, debug=False):
    """Convert a single page's Cloud Vision data to sorted Markdown text.
    
    Uses paragraph-level coordinates for better layout reconstruction.
    """
    blocks = page_data.get('blocks', [])
    
    # Step 1: Extract all paragraphs with coordinates
    all_paragraphs = []
    
    for block in blocks:
        # Skip non-text blocks
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        
        for para in block.get('paragraphs', []):
            text = extract_paragraph_text(para)
            text_stripped = text.strip()
            
            if not text_stripped:
                continue
            
            x_min, y_min, x_max, y_max = get_bounding_box(para)
            
            # Skip scanner artifacts
            if is_scanner_artifact(text_stripped):
                continue
            
            # Skip footer artifacts
            if is_footer_artifact(y_min):
                continue
            
            all_paragraphs.append((x_min, y_min, x_max, y_max, text_stripped))
    
    if not all_paragraphs:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n"
    
    # Step 2: Group into rows
    rows = group_paragraphs_into_rows(all_paragraphs)
    
    # Step 3: Output each row
    output_lines = [f"## PAGE {page_num}", ""]
    
    for row in rows:
        # Build row text: join paragraphs left-to-right
        row_texts = []
        prev_x_max = 0
        
        for i, (x_min, y_min, x_max, y_max, text) in enumerate(row):
            # Determine spacing between columns
            if i > 0:
                gap = x_min - prev_x_max
                if gap > 0.15:  # Large gap → column separator
                    row_texts.append('    ')  # 4-space column separator
                elif gap > 0.03:
                    row_texts.append('  ')  # Normal space
                else:
                    row_texts.append(' ')  # Tight
            
            row_texts.append(text)
            prev_x_max = x_max
        
        line = ''.join(row_texts)
        output_lines.append(line)
    
    output_lines.append("")
    
    if debug:
        # Add debug info: row count
        output_lines.insert(1, f"<!-- {len(rows)} rows, {len(all_paragraphs)} paragraphs -->")
    
    return '\n'.join(output_lines)


def json_to_markdown(json_path, output_dir, start_page, debug=False):
    """Process one JSON file → one 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 = []
    responses = data.get('responses', [])
    
    for i, response in enumerate(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:
            # Fallback: use raw text
            raw_text = fta.get('text', '').strip()
            if raw_text:
                page_text = f"## PAGE {page_num}\n\n{raw_text}\n"
            else:
                page_text = f"## PAGE {page_num}\n\n*(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(responses)


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v2.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:
            # Preview first 15 lines
            with open(out_path, 'r') as f:
                preview = ''.join(f.readline() for _ in range(15))
            print(f"\n{'='*60}")
            print(f"📄 {os.path.basename(out_path)} ({pages} trang) — PREVIEW:")
            print(f"{'='*60}")
            print(preview)
            print(f"... (dry-run, file 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()