#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V3 — Split multi-line blocks).
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.
Đặc biệt: tách các block nhiều dòng thành từng dòng riêng để xử lý TOC.

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

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),
    re.compile(r'CS\s*CamScanner', re.IGNORECASE),
]

# Header labels that appear as column headers in TOC (not real content)
HEADER_LABELS_MY = [
    re.compile(r'^စာမျက်နှာ\s*$'),
    re.compile(r'^မာတိကာခေါင်းစဉ်\s*$'),
    re.compile(r'^အကြောင်းအရာ\s*$'),
]

# Max Y for header zone
HEADER_Y_MAX = 0.12
# Min Y for footer zone  
FOOTER_Y_MIN = 0.88

# Y overlap threshold for row grouping
Y_OVERLAP_THRESHOLD = 0.15


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

def extract_paragraph_text_with_breaks(para):
    """Extract text from a paragraph, tracking line breaks.
    Returns (full_text, line_texts) where line_texts is list of lines.
    """
    text = ''
    lines = []
    current_line = ''
    
    for word in para.get('words', []):
        for sym in word.get('symbols', []):
            char = sym.get('text', '')
            text += char
            current_line += char
            
            prop = sym.get('property', {})
            if prop.get('detectedBreak'):
                bt = prop['detectedBreak']['type']
                if bt in ('SPACE', 'SURE_SPACE'):
                    text += ' '
                    current_line += ' '
                elif bt in ('EOL_SURE_SPACE', 'LINE_BREAK'):
                    text += '\n'
                    if current_line.strip():
                        lines.append(current_line.strip())
                    current_line = ''
    
    # Don't forget the last line
    if current_line.strip():
        lines.append(current_line.strip())
    
    return text, lines


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."""
    for pat in SCANNER_PATTERNS:
        if pat.search(text):
            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=Y_OVERLAP_THRESHOLD):
    """Check if two Y ranges overlap significantly."""
    overlap_start = max(y1_min, y2_min)
    overlap_end = min(y1_max, y2_max)
    overlap = max(0, overlap_end - overlap_start)
    
    range1 = max(y1_max - y1_min, 0.001)
    range2 = max(y2_max - y2_min, 0.001)
    min_range = min(range1, range2)
    
    return overlap / min_range


def split_multi_line_block(paragraphs):
    """Split multi-line paragraphs into individual line entries.
    
    For paragraphs with multiple lines (containing \n in text), 
    estimate Y position of each line and create separate entries.
    
    Returns list of (x_min, x_mid, x_max, text_line) for each line fragment.
    """
    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 len(lines) <= 1:
            # Single line: use as-is
            result.append((x_min, y_min, x_max, y_max, lines[0] if lines else text))
        else:
            # Multi-line: distribute lines evenly across Y range
            num_lines = len(lines)
            line_height = (y_max - y_min) / num_lines
            
            for i, line in enumerate(lines):
                line_y_min = y_min + i * line_height
                line_y_max = line_y_min + line_height
                result.append((x_min, line_y_min, x_max, line_y_max, line))
    
    return result


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), then x_min
    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
        
        overlap = y_overlap(current_y_min, current_y_max, y_min, y_max)
        
        if overlap > 0:
            current_row.append(para)
            current_y_min = min(current_y_min, y_min)
            current_y_max = max(current_y_max, y_max)
        else:
            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
    
    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."""
    blocks = page_data.get('blocks', [])
    
    # Step 1: Extract paragraphs with coordinates
    all_paragraphs = []
    
    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        
        for para in block.get('paragraphs', []):
            text, _ = extract_paragraph_text_with_breaks(para)
            text_stripped = text.strip()
            
            if not text_stripped:
                continue
            
            x_min, y_min, x_max, y_max = get_bounding_box(para)
            
            if is_scanner_artifact(text_stripped):
                continue
            
            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: Split multi-line blocks into individual lines
    line_fragments = split_multi_line_block(all_paragraphs)
    
    # Step 3: Group into rows
    rows = group_paragraphs_into_rows(line_fragments)
    
    # Step 4: Output each row
    output_lines = [f"## PAGE {page_num}", ""]
    
    for row in rows:
        row_texts = []
        prev_x_max = 0
        
        for i, (x_min, y_min, x_max, y_max, text) in enumerate(row):
            if i > 0:
                gap = x_min - prev_x_max
                if gap > 0.15:
                    row_texts.append('    ')
                elif gap > 0.03:
                    row_texts.append('  ')
                else:
                    row_texts.append(' ')
            
            row_texts.append(text)
            prev_x_max = x_max
        
        output_lines.append(''.join(row_texts))
    
    output_lines.append("")
    
    if debug:
        output_lines.insert(1, f"<!-- {len(rows)} rows, {len(line_fragments)} fragments -->")
    
    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:
            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_v3.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(20))
            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()