#!/usr/bin/env python3
"""
Cleanup OCR JSON → Markdown cho Pali-Taykha (v2 — dùng tọa độ)
- Xóa watermark "Scanned with CS CamScanner" (text-level)
- Xóa running headers dựa trên tọa độ (block-level)
- Phát hiện các block header lặp lại giữa các trang
"""
import json
import os
import glob
from collections import Counter

# ===== CONFIG =====
RAW_DIR = "/home/tuan-nguyen/.openclaw/workspace/010-pali-thaykha/ocr/raw"
OUT_DIR = "/home/tuan-nguyen/.openclaw/workspace/010-pali-thaykha/extracted"
HEADER_Y_MAX = 0.16     # Ngưỡng header zone (top of page)
Y_TOLERANCE = 0.012     # Tolerance khi so sánh tọa độ y
os.makedirs(OUT_DIR, exist_ok=True)

# ===== PASS 1: Quét tất cả block header, tìm pattern lặp =====
def extract_header_blocks(data):
    """Trích xuất tất cả block trong header zone, trả về dict {page: [blocks]}"""
    headers = {}
    for resp in data.get("responses", []):
        page = resp["context"]["pageNumber"]
        anno = resp.get("fullTextAnnotation", {})
        if not anno:
            headers[page] = []
            continue

        page_blocks = []
        for block in anno.get("pages", [{}])[0].get("blocks", []):
            bbox = block.get("boundingBox", {}).get("normalizedVertices", [{}]*4)
            ys = [v.get("y", 0) for v in bbox if "y" in v]
            xs = [v.get("x", 0) for v in bbox if "x" in v]
            if not ys or not xs:
                continue
            y_min, y_max = min(ys), max(ys)
            x_min, x_max = min(xs), max(xs)

            if y_max <= HEADER_Y_MAX:
                text_parts = []
                for para in block.get("paragraphs", []):
                    for word in para.get("words", []):
                        for symbol in word.get("symbols", []):
                            text_parts.append(symbol.get("text", ""))
                text = "".join(text_parts).strip()
                if text:
                    page_blocks.append({
                        "text": text,
                        "y_min": y_min,
                        "y_max": y_max,
                        "x_min": x_min,
                        "x_max": x_max,
                    })
        headers[page] = page_blocks
    return headers


def identify_running_headers(all_header_blocks):
    """Xác định text nào là running header (xuất hiện ≥ 2 trang)"""
    text_counter = Counter()
    text_pages = {}
    for page, blocks in all_header_blocks.items():
        for b in blocks:
            t = b["text"]
            text_counter[t] += 1
            if t not in text_pages:
                text_pages[t] = []
            text_pages[t].append(page)

    running = set()
    for text, count in text_counter.items():
        if count >= 2:
            # Running header: xuất hiện ở nhiều trang
            running.add(text)

    # Thêm: single-char artifacts luôn là header noise
    for text in text_counter:
        if len(text) == 1:
            running.add(text)

    # Thêm: page numbers (Myanmar/Arabic digits only)
    import re
    for text in text_counter:
        if re.match(r'^[၀-၉ဝ0-9]+$', text):
            running.add(text)

    return running, text_pages


def clean_page_with_coords(page_num, blocks, header_blocks, running_texts):
    """
    Xóa các block header khỏi fullTextAnnotation.text
    bằng cách xóa dòng text tương ứng
    """
    anno = blocks  # fullTextAnnotation từ JSON

    # Build set of texts to remove from this page
    texts_to_remove = set()
    for hb in header_blocks:
        if hb["text"] in running_texts:
            texts_to_remove.add(hb["text"])

    # Get page text
    raw_text = anno.get("text", "")

    # Xóa watermark
    import re
    raw_text = re.sub(r"[\n\s]*Scanned with[\n\s]+CS CamScanner['']?", "", raw_text, flags=re.UNICODE)

    # Xóa từng header text khỏi page text (chỉ xóa khi nó là dòng riêng hoặc đầu dòng)
    lines = raw_text.split("\n")
    cleaned_lines = []
    for line in lines:
        stripped = line.strip()
        should_skip = False
        for t in texts_to_remove:
            if stripped == t:
                should_skip = True
                break
            # Nếu dòng chứa text header nhưng không phải exact match, bỏ qua
            # (tránh xóa nhầm nội dung chứa từ trùng)
        if not should_skip:
            cleaned_lines.append(line)

    # Xóa dòng trống liên tiếp (>2)
    result = []
    empty_count = 0
    for line in cleaned_lines:
        if line.strip():
            result.append(line)
            empty_count = 0
        else:
            empty_count += 1
            if empty_count <= 2:
                result.append(line)

    # Trim đầu/cuối
    while result and not result[0].strip():
        result.pop(0)
    while result and not result[-1].strip():
        result.pop()

    return "\n".join(result)


def main():
    json_files = sorted(glob.glob(os.path.join(RAW_DIR, "*.json")))
    if not json_files:
        print("❌ No JSON files found!")
        return

    # === PASS 1: Quét tất cả để tìm running headers ===
    print("🔍 PASS 1: Quét header zone tất cả trang...")
    all_headers = {}
    all_data = {}
    for jf in json_files:
        with open(jf, "r", encoding="utf-8") as f:
            data = json.load(f)
        all_data[jf] = data
        h = extract_header_blocks(data)
        all_headers.update(h)

    running_texts, text_pages = identify_running_headers(all_headers)
    print(f"   Phát hiện {len(running_texts)} running header text:")
    for t in sorted(running_texts):
        pages = text_pages[t]
        print(f"     {repr(t):30s} → {len(pages)} pages: {pages[:5]}{'...' if len(pages)>5 else ''}")
    print()

    # === PASS 2: Cleanup từng file ===
    print("🔧 PASS 2: Cleanup + Export Markdown...")
    total_pages = 0
    total_headers_removed = 0
    total_watermarks = 0

    all_md_parts = []
    all_md_parts.append("# သောဠသမကျမ်း — OCR Cleaned v2\n\n")

    for jf in json_files:
        data = all_data[jf]
        basename = os.path.splitext(os.path.basename(jf))[0]
        md_path = os.path.join(OUT_DIR, f"{basename}.md")

        md_parts = []
        for resp in data.get("responses", []):
            if "fullTextAnnotation" not in resp:
                continue

            page_num = resp["context"]["pageNumber"]
            anno = resp["fullTextAnnotation"]

            # Lấy header blocks cho trang này
            page_headers = all_headers.get(page_num, [])
            header_count = sum(1 for hb in page_headers if hb["text"] in running_texts)

            cleaned = clean_page_with_coords(page_num, anno, page_headers, running_texts)

            if "Scanned with" in anno.get("text", ""):
                total_watermarks += 1

            total_pages += 1
            total_headers_removed += header_count

            md_parts.append(f"<!-- page {page_num} -->\n\n{cleaned}\n")

        full_md = "\n".join(md_parts)
        with open(md_path, "w", encoding="utf-8") as f:
            f.write(full_md)

        fsize = os.path.getsize(md_path)
        print(f"   ✅ {os.path.basename(jf):25s} → {os.path.basename(md_path):25s} "
              f"({fsize:,} bytes)")

        all_md_parts.append(full_md)

    # Merge
    merged_path = os.path.join(OUT_DIR, "Pali-Taykha-full.md")
    with open(merged_path, "w", encoding="utf-8") as out:
        out.write("\n".join(all_md_parts))
        out.write(f"\n\n---\n**Tổng:** {total_pages} trang | "
                  f"Headers đã xóa: {total_headers_removed} | "
                  f"Watermarks đã xóa: {total_watermarks}\n")

    print(f"\n{'='*60}")
    print(f"📊 Tổng: {total_pages} trang")
    print(f"   🗑️  Headers đã xóa: {total_headers_removed}")
    print(f"   🗑️  Watermarks đã xóa: {total_watermarks}")
    print(f"📄 Merged: {merged_path} ({os.path.getsize(merged_path):,} bytes)")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()