#!/usr/bin/env python3
"""
Cleanup OCR → Markdown cho Pali-Taykha (v4 — xóa theo tọa độ Y của số trang)
Nguyên lý: Phát hiện số trang ở header → xóa TOÀN BỘ text cùng dải tọa độ Y
"""
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.18      # vùng header (trên cùng trang)
Y_TOLERANCE = 0.012       # overlap tolerance
os.makedirs(OUT_DIR, exist_ok=True)

# Pattern số trang: Myanmar hoặc Ả Rập, 1-3 chữ số
PAGE_NUM_RE = re.compile(r'^[၀-၉ဝ0-9]{1,3}$')


def get_block_text(block):
    """Trích xuất text từ Cloud Vision 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):
    """Lấy bounding box của 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 find_page_number_y_range(blocks):
    """
    Tìm số trang trong header zone và trả về dải Y của header.
    Tìm block nhỏ (text ngắn, là số) ở rìa trái/phải vùng header.
    Nếu không tìm thấy → fallback: dải Y của tất cả block header.
    """
    header_blocks = []
    for block in blocks:
        bbox = get_block_bbox(block)
        if not bbox:
            continue
        y_min, y_max, x_min, x_max = bbox
        if y_max > HEADER_Y_MAX:
            continue
        text = get_block_text(block)
        if not text:
            continue
        header_blocks.append((y_min, y_max, x_min, x_max, text))

    if not header_blocks:
        return None

    # Tìm page number candidates: số ngắn, ở rìa trái (x<0.35) hoặc phải (x>0.65)
    page_num_ys = []
    for y_min, y_max, x_min, x_max, text in header_blocks:
        is_page_num = PAGE_NUM_RE.match(text) and (x_max < 0.35 or x_min > 0.65)
        is_artifact = (len(text) <= 3 and re.match(r'^[A-Za-z0-9၀-၉ဝ]+$', text) 
                       and (x_max < 0.35 or x_min > 0.65))
        if is_page_num or is_artifact:
            page_num_ys.append((y_min, y_max))

    if not page_num_ys:
        # Fallback: dùng toàn bộ header blocks
        all_y_min = min(h[0] for h in header_blocks)
        all_y_max = max(h[1] for h in header_blocks)
        return (all_y_min - Y_TOLERANCE, all_y_max + Y_TOLERANCE)

    # Lấy dải Y của số trang
    pn_y_min = min(p[0] for p in page_num_ys)
    pn_y_max = max(p[1] for p in page_num_ys)

    # Mở rộng: tìm TẤT CẢ block trong header zone overlap với dải Y của số trang
    band_y_min = pn_y_min
    band_y_max = pn_y_max
    for y_min, y_max, x_min, x_max, text in header_blocks:
        # Check overlap
        if y_max >= pn_y_min - Y_TOLERANCE and y_min <= pn_y_max + Y_TOLERANCE:
            band_y_min = min(band_y_min, y_min)
            band_y_max = max(band_y_max, y_max)

    return (band_y_min - 0.005, band_y_max + 0.005)


def get_header_texts_for_page(blocks, header_y_range):
    """Lấy tất cả text từ các block nằm trong dải Y của header"""
    if not header_y_range:
        return set()
    y_lo, y_hi = header_y_range
    texts = set()
    for block in blocks:
        bbox = get_block_bbox(block)
        if not bbox:
            continue
        y_min, y_max = bbox[0], bbox[1]
        # Block nằm trong hoặc overlap với header band
        if y_max >= y_lo and y_min <= y_hi:
            text = get_block_text(block)
            if text:
                texts.add(text)
    return texts


def clean_page_text(raw_text, header_texts):
    """Xóa header text khỏi page text (chỉ ở dòng đầu)"""
    # Xóa watermark
    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

    # Chỉ quét 10 dòng đầu (header luôn nằm ở đầu trang)
    scan_limit = min(10, len(lines))
    keep = [True] * len(lines)

    for i in range(scan_limit):
        stripped = lines[i].strip()
        if not stripped:
            continue
        if stripped in header_texts:
            keep[i] = False
            removed += 1
            continue
        # Check: dòng có phải là PREFIX của 1 header text merge dài không
        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
        # Check single-char artifacts ở header vùng
        if len(stripped) == 1 and i < 5:
            if re.match(r'^[A-Za-z¦|=\-*+]+$', stripped):
                keep[i] = False
                removed += 1
                continue

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

    # Xóa dòng trống liên tiếp >2
    result = []
    empty = 0
    for line in cleaned_lines:
        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 = 0
    total_removed = 0
    total_watermarks = 0

    all_md_parts = ["# သောဠသမကျမ်း — OCR Cleaned v4\n\n"]

    print(f"🔧 Cleanup v4: xóa header theo tọa độ Y của 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", [])

            # Tìm dải Y của header dựa trên vị trí số trang
            header_y_range = find_page_number_y_range(blocks)
            header_texts = get_header_texts_for_page(blocks, header_y_range)

            cleaned, removed = clean_page_text(raw_text, header_texts)

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

            total_pages += 1
            total_removed += removed

            # Log chi tiết
            if header_texts:
                y_info = f"y=[{header_y_range[0]:.3f},{header_y_range[1]:.3f}]" if header_y_range else "y=N/A"
                print(f"   Page {page_num:2d}  {y_info}  xóa: {sorted(header_texts)}")

            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)
        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"📊 {total_pages} trang | 🗑️ {total_removed} dòng header | 🗑️ {total_watermarks} watermarks")
    print(f"📄 {merged_path} ({os.path.getsize(merged_path):,} bytes)")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()