#!/usr/bin/env python3
"""
Cleanup header/footer leaks từ JSON Cloud Vision → Markdown sạch.

Detects and removes:
  - Even pages: page number + book title + separator artifacts
  - Odd pages:  section/chapter running headers (Paṇṇāsaka, vagga titles, etc.)

Strategy: Block-level detection (y-coordinate + pattern) → Text-level prefix removal.

Usage:
  python3 cleanup_header_leaks.py <json_dir> <output_dir>
"""
import os, sys, json, glob, re

# ── Constants ────────────────────────────────────────────────────

# Myanmar digits + 'ဝ' (letter wa, U+101D) which OCR confuses with digit zero '၀' (U+1040)
MYANMAR_DIGITS = set('၀၁၂၃၄၅၆၇၈၉ဝ')
BOOK_TITLE_RE = re.compile(r'^အင်္ဂုတ္တိုရ်')  # Must START with book title (not contain)
SECTION_MARKERS = {'ပထမ', 'ဒုတိယ', 'တတိယ', 'စတုတ္ထ', 'ပဉ္စမ', 'ဆဋ္ဌ', 'သတ္တမ'}
ARTIFACT_MARKERS = {'-', '=', '*', 'J', '..', '.', '/', '။'}
ARTIFACT_RE = re.compile(r"^'?\d+'?$")  # Quoted numbers like '60', also bare digits

# Odd-page running headers: section/chapter titles that repeat on left-hand pages
ODD_HEADER_RES = [
    re.compile(r'ပဏ္ဏာသက'),                 # "Paṇṇāsaka" (group of 50)
    re.compile(r'^\d+[\-\s]+\S+ဝဂ်'),        # "1-Rūpādivagga", "18-Aparā..."
    re.compile(r'^\d+[\-\s]+\S+ပါဠိ'),       # "15-Aṭṭhānapāḷi"
    re.compile(r'^\d+[\-\s]+\S+သုတ်'),        # Sutta headings
    re.compile(r'ပတ္ထနာနှင့်\s*ပဋိညာဉ်'),    # Introduction header
]

HEADER_Y_CUTOFF = 0.12


def is_myanmar_page_number(text):
    text = text.strip()
    if not text or len(text) > 5:
        return False
    return all(c in MYANMAR_DIGITS for c in text)


def is_book_title(text):
    """Only matches if text STARTS with book title (not just contains it)."""
    return bool(BOOK_TITLE_RE.search(text.replace(' ', '')))


def is_section_marker(text):
    t = text.strip().strip('[](). ')
    return t in SECTION_MARKERS


def is_artifact(text):
    t = text.strip()
    return t in ARTIFACT_MARKERS or bool(ARTIFACT_RE.match(t))


def is_odd_page_header(text):
    """Check if text matches odd-page running header patterns."""
    compact = text.replace(' ', '')
    for pat in ODD_HEADER_RES:
        if pat.search(compact):
            return True
    return False


def extract_block_text(block):
    words = []
    for para in block.get('paragraphs', []):
        for word in para.get('words', []):
            w_text = ''.join(s.get('text', '') for s in word.get('symbols', []))
            words.append(w_text)
    return ' '.join(words)


def get_block_y_min(block):
    nv = block.get('boundingBox', {}).get('normalizedVertices', [])
    if nv:
        return min(v.get('y', 0) for v in nv)
    return 1.0


def get_header_texts_to_remove(page_data, page_num):
    """
    Identify header block texts to remove.
    
    Even pages: page number + book title + adjacent artifacts
    Odd pages:  section/chapter running headers + adjacent artifacts
    """
    blocks = page_data.get('blocks', [])
    texts_to_remove = []
    primary_indices = set()
    is_odd = page_num % 2 == 1
    
    # Pass 1: Identify primary header blocks (first 3 blocks)
    for i, block in enumerate(blocks):
        if i > 2:
            break
        if block.get('blockType') not in (None, 'TEXT', 'text', 'TEXT_BLOCK'):
            continue
        
        y_min = get_block_y_min(block)
        text_compact = extract_block_text(block).replace(' ', '')
        
        if y_min >= HEADER_Y_CUTOFF:
            continue
        
        # Even-page headers: page number, book title, section markers
        if is_book_title(text_compact) or is_myanmar_page_number(text_compact) or is_section_marker(text_compact):
            texts_to_remove.append(text_compact)
            primary_indices.add(i)
        
        # Odd-page headers: section/chapter titles
        elif is_odd and is_odd_page_header(text_compact):
            texts_to_remove.append(text_compact)
            primary_indices.add(i)
        
        # Artifacts at top (y < 0.10)
        elif is_artifact(text_compact) and y_min < 0.10:
            texts_to_remove.append(text_compact)
            primary_indices.add(i)
    
    # Pass 2: Adjacent artifacts (within first 3 blocks)
    if primary_indices:
        for i, block in enumerate(blocks):
            if i > 2:
                break
            if block.get('blockType') not in (None, 'TEXT', 'text', 'TEXT_BLOCK'):
                continue
            if i in primary_indices:
                continue
            
            text_compact = extract_block_text(block).replace(' ', '')
            
            if is_artifact(text_compact):
                adjacent = any(abs(i - pi) <= 1 for pi in primary_indices)
                if adjacent:
                    texts_to_remove.append(text_compact)
    
    return texts_to_remove


def clean_page_text(full_text, header_texts_to_remove, page_num):
    """
    Remove header lines from full text.
    
    Strategies:
    1. Exact/partial match against detected header block texts
    2. Prefix check: book title, page numbers, section markers (first 5 lines)
    3. Odd-page header patterns (first 3 lines on odd pages)
    """
    if not full_text:
        return full_text
    
    lines = full_text.split('\n')
    kept_lines = []
    removed_any = False
    is_odd = page_num % 2 == 1
    
    header_compacts = set(header_texts_to_remove) if header_texts_to_remove else set()
    
    for line_idx, line in enumerate(lines):
        line_stripped = line.strip()
        line_compact = line_stripped.replace(' ', '')
        
        if not line_compact:
            if not kept_lines or not removed_any:
                continue
            kept_lines.append(line)
            continue
        
        # Strategy 1: Exact match
        if line_compact in header_compacts:
            removed_any = True
            continue
        
        # Strategy 1b: Partial match for long texts (common prefix >= 10 chars)
        matched = False
        for h in header_compacts:
            if len(h) > 5:
                common = min(len(line_compact), len(h))
                if common >= 10 and line_compact[:common] == h[:common]:
                    matched = True
                    removed_any = True
                    break
        if matched:
            continue
        
        # Strategy 2: Prefix-based detection (early lines only, <= 5)
        if line_idx < 5:
            # Book title prefix
            if line_compact.startswith('အင်္ဂုတ္တိုရ်ပါဠိတော်နိဿယ'):
                removed_any = True
                continue
            # Standalone Myanmar page number
            if is_myanmar_page_number(line_compact) and len(line_compact) <= 5:
                removed_any = True
                continue
            # Standalone section marker
            if is_section_marker(line_compact):
                removed_any = True
                continue
            # Artifact (quoted numbers, dashes, etc.)
            if is_artifact(line_compact):
                removed_any = True
                continue
        
        # Strategy 3: Odd-page header patterns (first 3 lines on odd pages)
        if is_odd and line_idx < 3 and is_odd_page_header(line_compact):
            removed_any = True
            continue
        
        kept_lines.append(line)
    
    # Trim leading/trailing empty lines
    while kept_lines and not kept_lines[0].strip():
        kept_lines.pop(0)
    while kept_lines and not kept_lines[-1].strip():
        kept_lines.pop()
    
    return '\n'.join(kept_lines)


def process_json_file(json_path, output_dir, start_page):
    """Process one JSON file → one clean Markdown file."""
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    basename = os.path.splitext(os.path.basename(json_path))[0]
    out_path = os.path.join(output_dir, f"{basename}.md")
    
    lines = []
    total_removed = 0
    
    for i, response in enumerate(data.get('responses', [])):
        page_num = start_page + i
        fta = response.get('fullTextAnnotation', {})
        pages = fta.get('pages', [])
        full_text = fta.get('text', '').strip()
        
        if pages and full_text:
            header_texts = get_header_texts_to_remove(pages[0], page_num)
            total_removed += len(header_texts)
            clean_text = clean_page_text(full_text, header_texts, page_num)
        else:
            clean_text = full_text
        
        lines.append(f"## PAGE {page_num}")
        lines.append("")
        if clean_text:
            lines.append(clean_text)
        else:
            lines.append("*(trang trống)*")
        lines.append("")
    
    os.makedirs(output_dir, exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(lines))
    
    return out_path, len(data.get('responses', [])), total_removed


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 cleanup_header_leaks.py <json_dir> <output_dir>")
        sys.exit(1)
    
    json_dir = sys.argv[1]
    output_dir = sys.argv[2]
    
    json_files = sorted(
        glob.glob(os.path.join(json_dir, '*.json')),
        key=lambda f: int(re.search(r'output-(\d+)-to', os.path.basename(f)).group(1))
    )
    
    if not json_files:
        print(f"❌ Không tìm thấy file JSON nào trong {json_dir}")
        sys.exit(1)
    
    print(f"📄 Cleanup header leaks từ {len(json_files)} file JSON...")
    print(f"   Ngưỡng: y < {HEADER_Y_CUTOFF} | Chẵn: page#+book | Lẻ: section titles")
    print()
    
    total_pages = 0
    total_removed = 0
    
    for jf in json_files:
        base = int(re.search(r'output-(\d+)-to', os.path.basename(jf)).group(1))
        out_path, pages, removed = process_json_file(jf, output_dir, base)
        total_pages += pages
        total_removed += removed
        
        tag = "✅" if removed > 0 else "▫️"
        print(f"   {tag} {os.path.basename(out_path)} ({pages} trang, xóa {removed})")
    
    print(f"\n🎉 Done! {len(json_files)} file .md → {os.path.abspath(output_dir)}")
    print(f"   Tổng: {total_pages} trang, xóa {total_removed} header blocks")


if __name__ == '__main__':
    main()