#!/usr/bin/env python3
"""
JSON → Markdown RAW v2: Parse theo tọa độ blocks thay vì fullTextAnnotation.text.

Khắc phục lỗi reading order của Google Vision với layout 2 cột.
Nguyên lý: Gom blocks thành "dòng" dựa trên Y-overlap, sắp xếp theo X trong mỗi dòng.

Usage:
  python3 json_to_markdown_v2_coords.py <json_path> <output_dir>
Example:
  python3 json_to_markdown_v2_coords.py "010-pali-thaykha/ocr/raw/output-7-to-9.json" "010-pali-thaykha/extracted-v2/"
"""
import os
import sys
import json
import re

# Ngưỡng Y-overlap để coi 2 blocks cùng dòng
Y_OVERLAP_THRESHOLD = 0.015  # 1.5% chiều cao trang


def get_block_bbox(block):
    """Trả về (x_min, y_min, x_max, y_max) từ normalizedVertices."""
    vertices = block.get('boundingBox', {}).get('normalizedVertices', [])
    if len(vertices) == 4:
        xs = [v.get('x', 0) for v in vertices]
        ys = [v.get('y', 0) for v in vertices]
        return (min(xs), min(ys), max(xs), max(ys))
    return (0, 0, 0, 0)


def get_block_text(block):
    """Trích text từ block (paragraphs → words → symbols)."""
    texts = []
    for para in block.get('paragraphs', []):
        words = []
        for word in para.get('words', []):
            symbols = ''.join(s.get('text', '') for s in word.get('symbols', []))
            words.append(symbols)
        texts.append(' '.join(words))
    return '\n'.join(texts)


def blocks_overlap_y(y1_min, y1_max, y2_min, y2_max):
    """Kiểm tra 2 blocks có overlap trục Y không (cùng dòng)."""
    overlap = min(y1_max, y2_max) - max(y1_min, y2_min)
    max_height = max(y1_max - y1_min, y2_max - y2_min)
    if max_height == 0:
        return False
    return overlap > 0 and (overlap / max_height) > 0.3


def group_blocks_into_lines(blocks_data):
    """
    Gom blocks thành các dòng dựa trên Y-overlap.
    Trong mỗi dòng, sắp xếp blocks theo X (trái → phải).
    """
    if not blocks_data:
        return []

    # Sort blocks by Y (top to bottom)
    sorted_blocks = sorted(blocks_data, key=lambda b: b['y_min'])

    lines = []
    current_line = [sorted_blocks[0]]
    current_y_min = sorted_blocks[0]['y_min']
    current_y_max = sorted_blocks[0]['y_max']

    for block in sorted_blocks[1:]:
        if blocks_overlap_y(current_y_min, current_y_max, block['y_min'], block['y_max']):
            # Same line
            current_line.append(block)
            current_y_min = min(current_y_min, block['y_min'])
            current_y_max = max(current_y_max, block['y_max'])
        else:
            # New line
            # Sort current line by X (left to right)
            current_line.sort(key=lambda b: b['x_min'])
            lines.append(current_line)
            current_line = [block]
            current_y_min = block['y_min']
            current_y_max = block['y_max']

    # Don't forget the last line
    if current_line:
        current_line.sort(key=lambda b: b['x_min'])
        lines.append(current_line)

    return lines


def blocks_to_markdown(blocks_data, page_num):
    """Chuyển blocks → Markdown text với thứ tự đúng theo tọa độ."""
    lines = group_blocks_into_lines(blocks_data)

    output = [f"## PAGE {page_num}", ""]

    for line_blocks in lines:
        # Nối text các blocks trong cùng dòng, cách nhau bởi space hoặc newline
        line_texts = []
        for b in line_blocks:
            text = b['text'].strip()
            if not text:
                continue

            # Nếu block chiếm full width (>70%) → text riêng dòng (heading/paragraph)
            block_width = b['x_max'] - b['x_min']
            if block_width > 0.7:
                # Full-width block: xuống dòng riêng
                if line_texts:
                    output.append(' '.join(line_texts))
                    line_texts = []
                output.append(text)
            else:
                line_texts.append(text)

        if line_texts:
            output.append(' '.join(line_texts))

    output.append("")
    return '\n'.join(output)


def process_json(json_path, output_dir):
    """Xử lý 1 file JSON → 1 file Markdown."""
    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}_v2.md")

    results = []
    for i, response in enumerate(data.get('responses', [])):
        fta = response.get('fullTextAnnotation', {})
        pages = fta.get('pages', [])
        if not pages:
            # Fallback: dùng text gốc
            text = fta.get('text', '')
            results.append(f"## PAGE {i + 1}\n\n{text}\n")
            continue

        # Parse blocks với tọa độ
        blocks_data = []
        for block in pages[0].get('blocks', []):
            x_min, y_min, x_max, y_max = get_block_bbox(block)
            text = get_block_text(block)
            if text.strip():
                blocks_data.append({
                    'x_min': x_min, 'y_min': y_min,
                    'x_max': x_max, 'y_max': y_max,
                    'text': text
                })

        # Sắp xếp lại theo tọa độ
        page_md = blocks_to_markdown(blocks_data, i + 1)
        results.append(page_md)

    os.makedirs(output_dir, exist_ok=True)
    final_text = '\n'.join(results)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write(final_text)

    return out_path


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 json_to_markdown_v2_coords.py <json_path_or_dir> <output_dir>")
        sys.exit(1)

    json_path = sys.argv[1]
    output_dir = sys.argv[2]

    if os.path.isdir(json_path):
        import glob
        json_files = sorted(glob.glob(os.path.join(json_path, '*.json')))
        for jf in json_files:
            out = process_json(jf, output_dir)
            print(f"✅ {os.path.basename(jf)} → {os.path.basename(out)}")
    else:
        out = process_json(json_path, output_dir)
        print(f"✅ {os.path.basename(json_path)} → {os.path.basename(out)}")

    print(f"\n📂 Output: {os.path.abspath(output_dir)}")


if __name__ == '__main__':
    main()