#!/usr/bin/env python3
"""
Cleanup OCR → Markdown cho Pali-Taykha (v5 — đơn giản & chắc chắn)
Nguyên lý: Nếu header zone có số trang → xóa TOÀN BỘ block trong zone đó
"""
import json, os, glob, re

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.15   # ngưỡng y_min: block bắt đầu trong vùng này → header
os.makedirs(OUT_DIR, exist_ok=True)

PAGE_NUM_RE = re.compile(r'^[၀-၉ဝ0-9]{1,3}$')


def get_block_text(block):
    parts = []
    for para in block.get("paragraphs", []):
        for word in para.get("words", []):
            for symbol in word.get("symbols", []):
                parts.append(symbol.get("text", ""))
    return "".join(parts).strip()


def get_block_bbox(block):
    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:
        return None
    return min(ys), max(ys), min(xs), max(xs)


def has_page_number(header_blocks):
    """Kiểm tra xem trong các block header có số trang không"""
    for y_min, y_max, x_min, x_max, text in header_blocks:
        if PAGE_NUM_RE.match(text):
            return True
        # Single-char artifact ở rìa (có thể là OCR lỗi của số trang)
        if len(text) <= 2 and (x_max < 0.30 or x_min > 0.70):
            if re.match(r'^[A-Za-z0-9၀-၉ဝ]+$', text):
                return True
    return False


def get_header_texts(blocks):
    """Lấy text từ TẤT CẢ block trong header zone (y_min < threshold)"""
    header_blocks = []
    footer_texts = set()
    for block in blocks:
        bbox = get_block_bbox(block)
        if not bbox:
            continue
        y_min, y_max, x_min, x_max = bbox
        text = get_block_text(block)
        if not text:
            continue
        # Header zone: y_min trong vùng header
        if y_min <= HEADER_Y_MAX:
            header_blocks.append((y_min, y_max, x_min, x_max, text))
        # Footer zone: thu thập text ở đáy trang (có thể bị Cloud Vision đẩy lên đầu)
        if y_min > 0.85:
            footer_texts.add(text)

    if not header_blocks:
        return set(), False

    # Chỉ xóa header nếu phát hiện số trang trong zone
    if not has_page_number(header_blocks):
        return footer_texts, False  # vẫn trả về footer texts để xóa

    header_texts = {text for _, _, _, _, text in header_blocks}
    header_texts.update(footer_texts)  # gộp footer
    return header_texts, True


def clean_page_text(raw_text, header_texts):
    """Xóa header text khỏi 10 dòng đầu"""
    text = re.sub(r"[\n\s]*Scanned with[\n\s]+CS CamScanner['']?", "", raw_text, flags=re.UNICODE)
    if not header_texts:
        return text, 0

    lines = text.split("\n")
    removed = 0
    scan_limit = min(10, len(lines))
    keep = [True] * len(lines)

    for i in range(scan_limit):
        stripped = lines[i].strip()
        if not stripped:
            continue

        # Exact match
        if stripped in header_texts:
            keep[i] = False
            removed += 1
            continue

        # Prefix match (xử lý block merge: header text bao gồm cả body)
        for ht in header_texts:
            if len(ht) > len(stripped) and ht.startswith(stripped):
                keep[i] = False
                removed += 1
                break
        if not keep[i]:
            continue

        # Single-char Latin artifacts ở 5 dòng đầu
        if len(stripped) == 1 and i < 5:
            if re.match(r'^[A-Za-z¦|=\-*+]+$', stripped):
                keep[i] = False
                removed += 1
                continue
        # Single digit artifact: nếu dòng trên đã bị xóa và dòng này là 1 ký tự số
        if i > 0 and len(stripped) == 1 and not keep[i-1]:
            if re.match(r'^[0-9]$', stripped):
                keep[i] = False
                removed += 1
                continue

    cleaned = [l for i, l in enumerate(lines) if keep[i]]

    # Compact empty lines
    result = []
    empty = 0
    for line in cleaned:
        if line.strip():
            result.append(line)
            empty = 0
        else:
            empty += 1
            if empty <= 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), removed


def main():
    json_files = sorted(glob.glob(os.path.join(RAW_DIR, "*.json")))
    total_pages = total_removed = total_watermarks = pages_with_headers = 0
    all_md = ["# သောဠသမကျမ်း — OCR Cleaned v5\n\n"]
    skipped = []

    print("🔧 Cleanup v5: xóa toàn bộ header zone nếu có số trang\n")

    for jf in json_files:
        with open(jf, "r", encoding="utf-8") as f:
            data = json.load(f)
        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", "")
            blocks = anno.get("pages", [{}])[0].get("blocks", [])

            header_texts, has_pn = get_header_texts(blocks)
            cleaned, removed = clean_page_text(raw_text, header_texts)

            if "Scanned with" in raw_text:
                total_watermarks += 1
            total_pages += 1
            total_removed += removed
            if has_pn:
                pages_with_headers += 1
            if not has_pn and header_texts:
                skipped.append((page_num, header_texts))

            if header_texts:
                print(f"   Page {page_num:2d}  {'✅' if has_pn else '⏭️ '}  xóa: {sorted(header_texts)}")

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

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

    merged = os.path.join(OUT_DIR, "Pali-Taykha-full.md")
    with open(merged, "w", encoding="utf-8") as out:
        out.write("\n".join(all_md))
        out.write(f"\n\n---\n**Tổng:** {total_pages} trang | "
                  f"Header: {pages_with_headers} trang có số trang | "
                  f"Đã xóa: {total_removed} dòng | Watermarks: {total_watermarks}\n")

    if skipped:
        print(f"\n⚠️  {len(skipped)} trang có block header nhưng không có số trang → bỏ qua:")
        for pn, texts in skipped:
            print(f"   Page {pn}: {sorted(texts)}")

    print(f"\n{'='*60}")
    print(f"📊 {total_pages} trang | {pages_with_headers} có header → xóa {total_removed} dòng")
    print(f"📄 {merged} ({os.path.getsize(merged):,} bytes)")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()