#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V4 — Column-aware layout).
- Phát hiện layout nhiều cột qua X-histogram
- Đọc column-by-column cho layout 2 cột độc lập
- Đọc row-by-row (trái→phải) cho layout TOC / 1 cột
- Tách multi-line blocks thành từng dòng riêng

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

import os
import sys
import json
import glob
import re
import math

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

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

# Column detection thresholds
MIN_COLUMN_GAP = 0.08       # Minimum X gap to consider as column boundary
MIN_COLUMN_WIDTH = 0.10     # Minimum width of a column cluster
MIN_BLOCKS_PER_COLUMN = 3   # Minimum blocks in a column to treat as separate column

# Y overlap threshold for row grouping
Y_OVERLAP_THRESHOLD = 0.15


# ─── Helpers ─────────────────────────────────────────────────────────────────

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):
    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):
    for pat in SCANNER_PATTERNS:
        if pat.search(text):
            return True
    return False


def y_overlap(y1_min, y1_max, y2_min, y2_max, threshold=Y_OVERLAP_THRESHOLD):
    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)
    return overlap / min(range1, range2)


# ─── Column Detection ────────────────────────────────────────────────────────

def detect_columns(paragraphs):
    """Detect if page has multi-column layout and find column boundaries.
    
    Uses X-midpoint histogram: looks for gaps in X distribution.
    Returns (num_columns, column_boundaries) where boundaries are 
    (x_min, x_max) for each column, or None if single column.
    """
    if len(paragraphs) < 6:
        return 1, [(0, 1)]
    
    # Get X midpoints, exclude very narrow items and full-width items
    x_data = []
    for x_min, y_min, x_max, y_max, text in paragraphs:
        width = x_max - x_min
        x_mid = (x_min + x_max) / 2
        # Skip items wider than 60% of page (likely full-width)
        if width > 0.60:
            continue
        x_data.append((x_mid, width, x_min, x_max, y_min, y_max, text))
    
    if len(x_data) < 6:
        return 1, [(0, 1)]
    
    # Sort by X midpoint
    x_data.sort(key=lambda d: d[0])
    
    # Find gaps between consecutive X midpoints
    # A gap is "significant" if the distance between items is large
    # AND there's no overlap in their X ranges
    
    gaps = []
    for i in range(len(x_data) - 1):
        x1_mid, w1, x1_min, x1_max, _, _, _ = x_data[i]
        x2_mid, w2, x2_min, x2_max, _, _, _ = x_data[i + 1]
        
        gap = x2_min - x1_max  # Actual gap between bounding boxes
        mid_gap = x2_mid - x1_mid
        
        if gap > MIN_COLUMN_GAP or (mid_gap > 0.15 and gap > 0.02):
            gaps.append((i, gap, mid_gap, x1_max, x2_min))
    
    if not gaps:
        return 1, [(0, 1)]
    
    # Find the most significant gap
    # Weight: prefer gaps with more items on each side
    best_gap = None
    best_score = 0
    
    for idx, gap, mid_gap, left_edge, right_edge in gaps:
        left_count = idx + 1
        right_count = len(x_data) - idx - 1
        
        if left_count >= MIN_BLOCKS_PER_COLUMN and right_count >= MIN_BLOCKS_PER_COLUMN:
            score = gap * 10 + min(left_count, right_count)
            if score > best_score:
                best_score = score
                best_gap = (left_edge, right_edge)
    
    if best_gap is None:
        return 1, [(0, 1)]
    
    # Calculate column boundaries
    boundary = (best_gap[0] + best_gap[1]) / 2
    
    col1_blocks = [d for d in x_data if d[3] < boundary]  # x_max < boundary
    col2_blocks = [d for d in x_data if d[2] > boundary]  # x_min > boundary
    
    # Only accept if both columns have enough content
    if len(col1_blocks) < MIN_BLOCKS_PER_COLUMN or len(col2_blocks) < MIN_BLOCKS_PER_COLUMN:
        return 1, [(0, 1)]
    
    # Determine column X ranges
    col1_x_min = min(d[2] for d in col1_blocks)
    col1_x_max = max(d[3] for d in col1_blocks)
    col2_x_min = min(d[2] for d in col2_blocks)
    col2_x_max = max(d[3] for d in col2_blocks)
    
    return 2, [(col1_x_min, col1_x_max), (col2_x_min, col2_x_max)]


# ─── Layout Processing ───────────────────────────────────────────────────────

def split_multi_line_paragraphs(paragraphs):
    """Split paragraphs with \n into individual line entries with estimated Y."""
    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:
            result.append((x_min, y_min, x_max, y_max, lines[0] if lines else text))
        else:
            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_into_rows(fragments):
    """Group fragments into rows based on Y overlap."""
    if not fragments:
        return []
    
    sorted_frags = sorted(fragments, key=lambda f: (f[1], f[0]))
    
    rows = []
    current_row = [sorted_frags[0]]
    current_y_min = sorted_frags[0][1]
    current_y_max = sorted_frags[0][3]
    
    for frag in sorted_frags[1:]:
        _, y_min, _, y_max, _ = frag
        overlap = y_overlap(current_y_min, current_y_max, y_min, y_max)
        
        if overlap > 0:
            current_row.append(frag)
            current_y_min = min(current_y_min, y_min)
            current_y_max = max(current_y_max, y_max)
        else:
            current_row.sort(key=lambda f: f[0])
            rows.append(current_row)
            current_row = [frag]
            current_y_min = y_min
            current_y_max = y_max
    
    if current_row:
        current_row.sort(key=lambda f: f[0])
        rows.append(current_row)
    
    return rows


def rows_to_text(rows, gap_threshold=0.15):
    """Convert rows to text with column separators."""
    lines = []
    for row in rows:
        parts = []
        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 > gap_threshold:
                    parts.append('    ')
                elif gap > 0.03:
                    parts.append('  ')
                else:
                    parts.append(' ')
            parts.append(text)
            prev_x_max = x_max
        lines.append(''.join(parts))
    return lines


def page_to_text(page_data, page_num, debug=False):
    """Convert a single page to sorted Markdown text with column detection."""
    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(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
            
            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
    fragments = split_multi_line_paragraphs(all_paragraphs)
    
    # Step 3: Detect column layout
    num_cols, col_bounds = detect_columns(all_paragraphs)
    
    output_lines = [f"## PAGE {page_num}", ""]
    
    if num_cols == 1:
        # Single column: row-by-row left-to-right
        rows = group_into_rows(fragments)
        output_lines.extend(rows_to_text(rows))
    else:
        # Multi-column: read column-by-column
        if debug:
            output_lines.append(f"<!-- {num_cols}-column layout detected, boundaries: {col_bounds} -->")
        
        for col_idx, (col_x_min, col_x_max) in enumerate(col_bounds):
            # Get fragments that belong to this column
            col_fragments = []
            for frag in fragments:
                x_min, y_min, x_max, y_max, text = frag
                x_mid = (x_min + x_max) / 2
                # Fragment belongs to column if its midpoint is within range
                # OR if it's narrow and contained within column bounds
                if col_x_min <= x_mid <= col_x_max:
                    col_fragments.append(frag)
                elif x_max - x_min < 0.25 and (x_min >= col_x_min - 0.05 and x_max <= col_x_max + 0.05):
                    col_fragments.append(frag)
            
            # Sort column fragments by Y
            col_fragments.sort(key=lambda f: (f[1], f[0]))
            
            if col_idx > 0:
                output_lines.append("")  # Blank line between columns
            
            # Group column fragments into rows (within column)
            col_rows = group_into_rows(col_fragments)
            
            # For multi-column: use smaller gap threshold since within-column items are close
            col_lines = rows_to_text(col_rows, gap_threshold=0.05)
            output_lines.extend(col_lines)
    
    output_lines.append("")
    
    if debug:
        # Add stats
        total_frags = len(fragments)
        output_lines.insert(1, f"<!-- {total_frags} fragments, {num_cols} column(s) -->")
    
    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 = []
    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()
            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(responses)


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v4.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()