#!/usr/bin/env python3
"""
Cleanup OCR JSON → Markdown cho Pali-Taykha (သောဠသမကျမ်း)
Pattern đặc thù:
  - Watermark "Scanned with\nCS CamScanner" ở cuối mọi trang
  - Page 1 có artifact "!! -" và publisher info
  - Không có running header như Tăng Chi Bộ
"""
import json
import os
import glob
import re

# ===== CONFIG =====
RAW_DIR = "/home/tuan-nguyen/.openclaw/workspace/010-pali-thaykha/ocr/raw"
OUT_DIR = "/home/tuan-nguyen/.openclaw/workspace/010-pali-thaykha/extracted"
os.makedirs(OUT_DIR, exist_ok=True)

# Pattern cần xóa
WATERMARK_RE = re.compile(r"[\n\s]*Scanned with[\n\s]+CS CamScanner['']?", re.UNICODE)
ARTIFACT_RE = re.compile(r"^!!\s*-?$", re.MULTILINE)

def clean_page_text(text, page_number):
    """Làm sạch text của một trang."""
    original = text

    # 1. Xóa watermark "Scanned with CS CamScanner"
    text = WATERMARK_RE.sub("", text)

    # 2. Page 1 specific: xóa artifact "!! -"
    if page_number == 1:
        text = ARTIFACT_RE.sub("", text)

    # 3. Chuẩn hóa whitespace (giữ cấu trúc dòng)
    # Xóa dòng trống liên tiếp (>2)
    lines = text.split("\n")
    cleaned_lines = []
    empty_count = 0
    for line in lines:
        stripped = line.strip()
        if stripped:
            cleaned_lines.append(line)
            empty_count = 0
        else:
            empty_count += 1
            if empty_count <= 2:
                cleaned_lines.append(line)

    # Bỏ các dòng trống ở đầu/cuối
    while cleaned_lines and not cleaned_lines[0].strip():
        cleaned_lines.pop(0)
    while cleaned_lines and not cleaned_lines[-1].strip():
        cleaned_lines.pop()

    return "\n".join(cleaned_lines)


def process_json_file(json_path):
    """Xử lý 1 file JSON → output markdown"""
    basename = os.path.splitext(os.path.basename(json_path))[0]
    md_path = os.path.join(OUT_DIR, f"{basename}.md")

    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    md_parts = []
    stats = {"pages": 0, "watermarks_removed": 0}

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

        page_num = resp.get("context", {}).get("pageNumber", "?")
        raw_text = resp["fullTextAnnotation"].get("text", "")

        cleaned = clean_page_text(raw_text, page_num)

        # Đếm watermark đã xóa
        if "Scanned with" in raw_text:
            stats["watermarks_removed"] += 1

        stats["pages"] += 1
        md_parts.append(f"<!-- page {page_num} -->\n\n{cleaned}\n")

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

    return stats, md_path


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

    print(f"🔧 Cleanup {len(json_files)} JSON files → Markdown")
    print(f"   Input:  {RAW_DIR}")
    print(f"   Output: {OUT_DIR}")
    print()

    total_pages = 0
    total_watermarks = 0
    total_chars = 0

    for jf in json_files:
        stats, md_path = process_json_file(jf)
        total_pages += stats["pages"]
        total_watermarks += stats["watermarks_removed"]

        fsize = os.path.getsize(md_path)
        total_chars += fsize
        fname = os.path.basename(jf)
        print(f"   ✅ {fname:25s} → {os.path.basename(md_path):25s} "
              f"({stats['pages']} pages, {fsize:,} bytes)")

    print(f"\n{'='*60}")
    print(f"📊 Tổng: {total_pages} trang, {total_watermarks} watermarks đã xóa")
    print(f"📁 Output: {total_chars:,} bytes ({total_chars/1024:.0f} KB)")

    # Gộp tất cả thành 1 file markdown duy nhất
    merged_path = os.path.join(OUT_DIR, "Pali-Taykha-full.md")
    all_md = sorted(glob.glob(os.path.join(OUT_DIR, "output-*.md")))
    with open(merged_path, "w", encoding="utf-8") as out:
        out.write("# သောဠသမကျမ်း — OCR Cleaned\n\n")
        out.write(f"**Tổng:** {total_pages} trang | **Watermarks đã xóa:** {total_watermarks}\n\n")
        out.write("---\n\n")
        for md in all_md:
            with open(md, "r", encoding="utf-8") as f:
                out.write(f.read())
            out.write("\n")

    merged_size = os.path.getsize(merged_path)
    print(f"📄 Merged: {merged_path} ({merged_size:,} bytes)")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()