#!/usr/bin/env python3
"""
clean_header_v1.py — Xóa header line đầu tiên mỗi trang nếu chứa số ở đầu/cuối
Áp dụng cho các batch từ output-67-to-69 trở về sau (page ≥ 67)

Quy tắc:
  - Xét dòng text đầu tiên của mỗi trang (sau ## PAGE X và annotation)
  - Nếu dòng bắt đầu HOẶC kết thúc bằng chữ số (Myanmar/Latin) → candidate header
  - Chỉ xóa nếu len(dòng) ≤ REFERENCE_LENGTH
  - Xóa luôn dòng trống theo sau header (nếu có)

Usage:
  python3 clean_header_v1.py <extracted_dir> [--dry-run]
"""

import os, sys, re, glob

# ── Config ──
REFERENCE_LENGTH = float('inf')  # Bỏ giới hạn độ dài — xóa mọi header có số ở đầu/cuối
MIN_PAGE = 67  # Chỉ áp dụng từ trang 67 trở đi
START_BATCH = "output-67-to-69"

# Regex patterns
MYANMAR_DIGITS = "၀၁၂၃၄၅၆၇၈၉"
DIGIT_START = re.compile(rf'^[{MYANMAR_DIGITS}0-9]')
DIGIT_END = re.compile(rf'[{MYANMAR_DIGITS}0-9]\s*$')


def _start_page(batch_name):
    """Extract start page from batch filename."""
    m = re.search(r'output-(\d+)-to-\d+', batch_name)
    return int(m.group(1)) if m else 0


def clean_page(text, page_num, dry_run=False):
    """
    Remove header line from page text if it matches header pattern.
    Returns (cleaned_text, was_cleaned, header_text).
    """
    lines = text.split("\n")
    
    # Find first non-empty, non-annotation line
    first_text_idx = None
    for i, line in enumerate(lines):
        stripped = line.strip()
        if not stripped:
            continue
        if stripped.startswith("<!--"):
            continue
        first_text_idx = i
        break
    
    if first_text_idx is None:
        return text, False, ""  # Empty page
    
    first_line = lines[first_text_idx].strip()
    
    # Check if it matches header pattern
    has_num_start = bool(DIGIT_START.match(first_line))
    has_num_end = bool(DIGIT_END.search(first_line))
    
    if not (has_num_start or has_num_end):
        return text, False, ""  # No number → not a header
    
    if len(first_line) > REFERENCE_LENGTH:
        return text, False, ""  # Too long → likely body text
    
    # It's a header! Remove it
    if dry_run:
        return text, True, first_line
    
    # Remove the header line
    del lines[first_text_idx]
    
    # Also remove trailing empty line right after header (if any)
    if first_text_idx < len(lines) and lines[first_text_idx].strip() == "":
        del lines[first_text_idx]
    
    return "\n".join(lines), True, first_line


def clean_file(filepath, dry_run=False):
    """Clean all pages in a markdown file."""
    with open(filepath) as f:
        content = f.read()
    
    # Split by page markers
    # Pattern: "## PAGE N\n<!-- ... -->\n..."
    pages = re.split(r'(## PAGE \d+)', content)
    
    # pages[0] is empty or preamble, then alternating: PAGE_MARKER, PAGE_CONTENT
    result_parts = []
    cleaned_count = 0
    total_pages = 0
    headers_found = []
    
    # pages[0] is empty (split before first marker), start from pages[1]
    i = 1
    
    while i < len(pages):
        if i + 1 >= len(pages):
            result_parts.append(pages[i])
            break
        
        page_marker = pages[i]
        page_content = pages[i + 1]
        
        # Extract page number
        m = re.match(r'## PAGE (\d+)', page_marker)
        page_num = int(m.group(1)) if m else 0
        total_pages += 1
        
        if page_num < MIN_PAGE:
            # Skip — before page 67
            result_parts.append(page_marker)
            result_parts.append(page_content)
        else:
            cleaned_content, was_cleaned, header = clean_page(page_content, page_num, dry_run)
            if was_cleaned:
                cleaned_count += 1
                headers_found.append(f"  PAGE {page_num}: '{header}'")
            result_parts.append(page_marker)
            result_parts.append(cleaned_content)
        
        i += 2
    
    new_content = "".join(result_parts)
    
    if not dry_run and cleaned_count > 0:
        with open(filepath, "w") as f:
            f.write(new_content)
    
    return cleaned_count, total_pages, headers_found


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 clean_header_v1.py <extracted_dir> [--dry-run]")
        sys.exit(1)
    
    extracted_dir = sys.argv[1]
    dry_run = "--dry-run" in sys.argv
    
    print(f"🔧 clean_header_v1 — Xóa header dòng đầu có số ở đầu/cuối")
    print(f"   REFERENCE: none (bỏ giới hạn độ dài)")
    print(f"   MIN_PAGE: {MIN_PAGE}\n")
    
    # Get all markdown files, sort naturally by start page
    md_files = glob.glob(os.path.join(extracted_dir, "output-*.md"))
    md_files.sort(key=lambda f: _start_page(os.path.basename(f)))
    
    # Filter: only batches starting at or after MIN_PAGE
    target_files = [f for f in md_files if _start_page(os.path.basename(f)) >= MIN_PAGE]
    
    if not target_files:
        print(f"❌ Không tìm thấy file nào từ page {MIN_PAGE} trở đi")
        sys.exit(1)
    
    print(f"📄 {len(target_files)} batch cần xử lý (từ {os.path.basename(target_files[0])})\n")
    
    total_cleaned = 0
    total_pages = 0
    all_headers = []
    
    for fpath in target_files:
        fname = os.path.basename(fpath)
        cleaned, pages, headers = clean_file(fpath, dry_run)
        total_cleaned += cleaned
        total_pages += pages
        if cleaned > 0:
            print(f"   ✅ {fname}: {cleaned}/{pages} pages cleaned")
            for h in headers:
                print(h)
        all_headers.extend(headers)
    
    mode = "DRY-RUN" if dry_run else "APPLIED"
    print(f"\n{'='*60}")
    print(f"📊 [{mode}] {total_cleaned}/{total_pages} pages có header → đã xóa")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()