#!/usr/bin/env python3
"""
Cleanup OCR JSON → Markdown cho Pali-Taykha (v3 — text-level header removal)
- PASS 1: Quét block header zone (y < 0.16) → tìm running header texts
- PASS 2: Xóa các dòng đầu trang khớp với running header patterns
- Xử lý được block merge (Cloud Vision gộp header + body)
"""
import json
import os
import glob
import re
from collections import Counter

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
os.makedirs(OUT_DIR, exist_ok=True)


def extract_header_texts(data):
    """PASS 1: Quét tất cả block trong header zone, trả về set các text"""
    header_texts = Counter()
    for resp in data.get("responses", []):
        anno = resp.get("fullTextAnnotation", {})
        if not anno:
            continue
        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]
            if not ys or max(ys) > HEADER_Y_MAX:
                continue
            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:
                header_texts[text] += 1
    return header_texts


def identify_running_headers(all_header_texts):
    """Xác định text nào là running header"""
    running = set()
    for text, count in all_header_texts.items():
        if count >= 2:
            running.add(text)

    # Thêm: single-char artifacts (OCR noise)
    for text in all_header_texts:
        if len(text) == 1:
            running.add(text)

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

    return running


def clean_page_text(raw_text, running_headers, page_num):
    """PASS 2: Xóa running header lines khỏi page text"""
    # Xóa watermark
    text = re.sub(r"[\n\s]*Scanned with[\n\s]+CS CamScanner['']?", "", raw_text, flags=re.UNICODE)

    lines = text.split("\n")

    # Xác định vùng đầu trang cần quét header (tối đa 8 dòng đầu)
    scan_limit = min(8, len(lines))

    # Đánh dấu dòng cần xóa
    to_remove = set()
    for i in range(scan_limit):
        stripped = lines[i].strip()
        if not stripped:
            continue

        # Check 1: exact match với running header
        if stripped in running_headers:
            to_remove.add(i)
            continue

        # Check 2: dòng là page number (Myanmar/Arabic digits)
        if re.match(r'^[၀-၉ဝ0-9]{1,3}$', stripped):
            to_remove.add(i)
            continue

        # Check 3: single-char artifact (nhưng không phải dấu câu hợp lệ)
        if len(stripped) == 1 and not re.match(r'[။၊\-\)]', stripped):
            # Chỉ xóa nếu dòng đó KHÔNG phải là ký tự Myanmar/Pali hợp lệ đứng một mình
            # (tránh xóa nhầm chữ cái thật)
            if not re.match(r'[\u1000-\u109F]', stripped):
                to_remove.add(i)
                continue

    # Xóa các dòng đã đánh dấu
    cleaned_lines = [l for i, l in enumerate(lines) if i not in to_remove]

    # Xóa dòng trống liên tiếp (>2) và trim
    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)

    while result and not result[0].strip():
        result.pop(0)
    while result and not result[-1].strip():
        result.pop()

    return "\n".join(result), len(to_remove)


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 (y < 0.16) tất cả trang...")
    all_header_texts = Counter()
    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_texts(data)
        all_header_texts.update(h)

    running_headers = identify_running_headers(all_header_texts)
    print(f"   Phát hiện {len(running_headers)} running header patterns:")
    for t in sorted(running_headers):
        count = all_header_texts.get(t, 0)
        if count >= 2:
            print(f"     {repr(t):30s} x{count}")
    print(f"   + {sum(1 for t in running_headers if all_header_texts.get(t,0) < 2)} single-char/page-number artifacts")
    print()

    # === PASS 2: Cleanup từng file ===
    print("🔧 PASS 2: Text-level cleanup + Export Markdown...")
    total_pages = 0
    total_removed = 0
    total_watermarks = 0

    all_md_parts = ["# သောဠသမကျမ်း — OCR Cleaned v3\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"]
            raw_text = anno.get("text", "")

            cleaned, removed = clean_page_text(raw_text, running_headers, page_num)

            if "Scanned with" in raw_text:
                total_watermarks += 1

            total_pages += 1
            total_removed += removed

            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"Dòng header đã xóa: {total_removed} | "
                  f"Watermarks: {total_watermarks}\n")

    print(f"\n{'='*60}")
    print(f"📊 Tổng: {total_pages} trang")
    print(f"   🗑️  Dòng header đã xóa: {total_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()