#!/usr/bin/env python3
"""
Clean leftover OCR headers/page numbers from Sadi final files (001-095).

Logic:
  - Repeating chapter names at page tops → REMOVE
  - Standalone page numbers (Myanmar/Arabic) near chapter names → REMOVE
  - Combined chapter+page on same line → REMOVE (keep body text after)
  - Page markers like "### Page X (YYY)" → REMOVE
  - "ရုပ်ပုံ ရှင်ကျင့်ဝတ်" with page numbers → REMOVE
"""

import os
import re
import shutil
from pathlib import Path
from collections import Counter

FINAL_DIR = Path("/home/tuan-nguyen/.openclaw/workspace/Chuan Muc Sadi/extracted/final")
BACKUP_DIR = FINAL_DIR / "backup"

MYANMAR_DIGITS = "၀-၉"
ARABIC_DIGITS = "0-9"
ALL_DIGITS = f"[{MYANMAR_DIGITS}{ARABIC_DIGITS}]"  # combined character class

# --- Chapter header patterns (repeating across pages) ---
# These appear at the top of each page as running headers
CHAPTER_PATTERNS = [
    r"သေခိယသိက္ခာပုဒ်များကို\s*သင်ပြခဏ်း",       # pure chapter name
    r"ခန္ဓကဝတ်ကို\s*သင်ပြခဏ်း",                    # pure chapter name  
]

# --- Combined chapter+page patterns ---
# e.g., "သေခိယသိက္ခာပုဒ်များကို သင်ပြခဏ်း ၆၁"
COMBINED_CHAPTER_PAGE = [
    (r"သေခိယသိက္ခာပုဒ်များကို\s*သင်ပြခဏ်း", "chapter"),
    (r"ခန္ဓကဝတ်ကို\s*သင်ပြခဏ်း", "chapter"),
]

# --- Other page header patterns ---
PAGE_HEADERS = [
    r"ရုပ်ပုံ\s*ရှင်ကျင့်ဝတ်",   # appears with page numbers as running header
]

# Compile chapter patterns
CHAPTER_RES = [re.compile(p) for p in CHAPTER_PATTERNS]

# --- Detection functions ---

def is_standalone_page_num(line: str) -> bool:
    """Check if line is a standalone page number (Myanmar 2-3 digits or Arabic 1-3 digits)."""
    stripped = line.strip()
    if not stripped:
        return False
    # Myanmar numerals: 2-3 digits
    if re.match(rf"^[{MYANMAR_DIGITS}]{{2,3}}$", stripped):
        return True
    # Arabic numerals: 1-3 digits
    if re.match(rf"^[{ARABIC_DIGITS}]{{1,3}}$", stripped):
        return True
    return False

def is_standalone_page_num_ext(line: str) -> bool:
    """Check if line is a standalone page number, including ဝ as zero digit (OCR error)."""
    stripped = line.strip()
    if not stripped:
        return False
    # Myanmar numerals with ဝ as zero: 2-3 digits
    if re.match(rf"^[{MYANMAR_DIGITS}ဝ]{{2,3}}$", stripped):
        return True
    # Arabic numerals: 1-3 digits
    if re.match(rf"^[{ARABIC_DIGITS}]{{1,3}}$", stripped):
        return True
    return False

def is_chapter_header(line: str) -> bool:
    """Check if line is a known repeating chapter header."""
    stripped = line.strip()
    if not stripped:
        return False
    for cre in CHAPTER_RES:
        if cre.fullmatch(stripped):
            return True
    return False

def is_combined_chapter_page(line: str) -> bool:
    """Check if line contains a chapter name followed by a page number.
    e.g., "သေခိယသိက္ခာပုဒ်များကို သင်ပြခဏ်း ၆၁"
    """
    stripped = line.strip()
    for pattern, _ in COMBINED_CHAPTER_PAGE:
        m = re.match(rf"^({pattern})\s+[{MYANMAR_DIGITS}{ARABIC_DIGITS}]{{1,3}}\s*$", stripped)
        if m:
            return True
    return False

def is_page_header_with_num(line: str) -> bool:
    """Check if a page header like 'ရုပ်ပုံ ရှင်ကျင့်ဝတ်' appears with page numbers nearby."""
    stripped = line.strip()
    if not stripped:
        return False
    for ph in PAGE_HEADERS:
        if re.fullmatch(ph, stripped):
            return True
    return False

def is_page_marker(line: str) -> bool:
    """Check if line is '### Page X (YYY)' style marker."""
    stripped = line.strip()
    return bool(re.match(rf"^#+\s*[Pp]age\s+[{ARABIC_DIGITS}]+\s*\([{MYANMAR_DIGITS}]+\)\s*$", stripped))


def clean_file(filepath: Path) -> tuple[int, list[str]]:
    """
    Clean a single file.
    Strategy: identify repeating chapter headers (appear 2+ times after page breaks),
    then remove them along with adjacent page numbers.
    """
    with open(filepath, "r", encoding="utf-8") as f:
        lines = f.readlines()

    n = len(lines)
    
    # --- Pass 1: Identify positions of chapter headers and page numbers ---
    # Store (line_index, type) for lines to potentially remove
    chapter_positions = []  # indices where chapter headers appear
    page_marker_positions = []  # indices of "### Page X" markers
    page_num_positions = []  # indices of standalone page numbers
    combined_positions = []  # indices of combined chapter+page lines
    page_header_positions = []  # indices of "ရုပ်ပုံ ရှင်ကျင့်ဝတ်"
    
    for i, line in enumerate(lines):
        stripped = line.strip()
        if not stripped:
            continue
        
        if is_chapter_header(stripped):
            chapter_positions.append(i)
        elif is_combined_chapter_page(stripped):
            combined_positions.append(i)
        elif is_page_marker(stripped):
            page_marker_positions.append(i)
        elif is_standalone_page_num(stripped):
            page_num_positions.append(i)
        elif is_page_header_with_num(stripped):
            page_header_positions.append(i)
    
    # --- Pass 2: Determine which lines to remove ---
    remove_set = set()
    
    # 2a: "### Page X (YYY)" markers → always remove
    remove_set.update(page_marker_positions)
    
    # 2b: Combined chapter+page on same line → always remove
    remove_set.update(combined_positions)
    
    # 2c: Chapter headers that appear 2+ times (repeating page headers)
    if len(chapter_positions) >= 2:
        remove_set.update(chapter_positions)
    elif len(chapter_positions) == 1:
        # Single appearance - only remove if it's right after ## Trang marker
        # and looks like a page header (not a section title)
        idx = chapter_positions[0]
        # Check if preceded by ## Trang marker within 3 lines
        is_after_page_break = False
        for j in range(max(0, idx - 3), idx):
            if re.match(r'^##\s+Trang', lines[j].strip()):
                is_after_page_break = True
                break
        if is_after_page_break:
            # Only remove if it appears right at the start of content (within 2 lines of ## Trang)
            lines_between = 0
            for j in range(idx - 1, -1, -1):
                stripped = lines[j].strip()
                if re.match(r'^##\s+Trang', stripped):
                    break
                if stripped:
                    lines_between += 1
            if lines_between <= 1:
                remove_set.add(idx)
    
    # 2d: Page numbers → remove if:
    #   - They are adjacent to (or within 1 line of) a chapter header
    #   - OR they are right after a ## Trang marker within 2 lines
    for idx in page_num_positions:
        should_remove = False
        
        # Check adjacency to chapter headers (already in remove_set or not)
        for cidx in chapter_positions:
            if abs(idx - cidx) <= 2:
                should_remove = True
                break
        
        if not should_remove:
            # Check adjacency to combined chapter+page lines
            for cidx in combined_positions:
                if abs(idx - cidx) <= 1:
                    should_remove = True
                    break
        
        if not should_remove:
            # Check if right after ## Trang marker (within 2 lines, with only empty/whitespace in between)
            for j in range(max(0, idx - 3), idx):
                stripped_j = lines[j].strip()
                if re.match(r'^##\s+Trang', stripped_j):
                    # Count non-empty lines between ## Trang and page number
                    gap_lines = 0
                    for k in range(j + 1, idx):
                        if lines[k].strip():
                            gap_lines += 1
                    if gap_lines <= 1:
                        should_remove = True
                    break
        
        if not should_remove:
            # Check adjacency to page header (ရုပ်ပုံ ရှင်ကျင့်ဝတ်)
            for hidx in page_header_positions:
                if abs(idx - hidx) <= 1:
                    should_remove = True
                    break
        
        if should_remove:
            remove_set.add(idx)
    
    # 2e: "ရုပ်ပုံ ရှင်ကျင့်ဝတ်" headers → remove if adjacent to page numbers
    for idx in page_header_positions:
        # Check adjacency to page numbers
        for pidx in page_num_positions:
            if abs(idx - pidx) <= 1:
                remove_set.add(idx)
                break
    
    # --- Pass 3: Build cleaned output ---
    cleaned = []
    removed_count = 0
    for i, line in enumerate(lines):
        if i in remove_set:
            removed_count += 1
            continue
        cleaned.append(line)
    
    # Write if changes were made
    if removed_count > 0:
        with open(filepath, "w", encoding="utf-8") as f:
            f.writelines(cleaned)
    
    return removed_count


def main():
    BACKUP_DIR.mkdir(exist_ok=True)

    # Target: Sadi-001 through Sadi-095
    files = sorted(
        [f for f in FINAL_DIR.glob("Sadi-0[0-9]*.md") 
         if int(re.match(r"Sadi-0*(\d+)", f.name).group(1)) < 96],
        key=lambda f: f.name,
    )

    if not files:
        print("No target files found (001-095).")
        return

    # First backup all
    for fpath in files:
        backup_path = BACKUP_DIR / fpath.name
        shutil.copy2(fpath, backup_path)

    print(f"Processing {len(files)} files (Sadi-001 to Sadi-095)...\n")

    total_removed = 0
    total_lines = 0

    for fpath in files:
        # Re-read after backup
        with open(fpath, "r", encoding="utf-8") as f:
            orig_lines = len(f.readlines())
        
        removed = clean_file(fpath)
        total_removed += removed
        total_lines += orig_lines

        status = f"✅ {removed} dòng bị xóa" if removed > 0 else "⬜ Sạch"
        print(f"  {fpath.name:40s} {status} ({orig_lines} → {orig_lines - removed} dòng)")

    print(f"\n{'='*60}")
    print(f"Tổng: {total_removed} dòng đã xóa / {total_lines} dòng ({len(files)} file)")
    print(f"Backup: {BACKUP_DIR}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()