#!/usr/bin/env python3
"""
Post-process Markdown RAW → clean Markdown cho Sổ tay Mahavihara.

Xóa:
  - Header running: မဟာဝိဟာရကျောင်းတိုက်(လှည်းကူး)
  - Header running: ဝတ်ရွတ်စဉ်
  - Single-char/digit artifacts ở đầu trang
  - Watermark "Scanned with ES CamScanner'"
  - Page numbers ở đầu trang

Usage:
  python3 cleanup_mahavihara.py <input_dir> <output_dir>
"""
import os, sys, glob, re
from pathlib import Path

# Watermark lines
WATERMARK_LINES = [
    'Scanned with',
    'ES CamScanner\'',
    "ES CamScanner'",
]

# Artifact patterns (đầu trang)
ARTIFACT_RE = re.compile(r'^[LJ\-=*/+."\':]$')           # single char artifacts
ARTIFACT_NUM_RE = re.compile(r'^[၀-၉ဝ]{1,3}$')          # Myanmar page numbers: "၁", "၁၄"
ARTIFACT_NUM_PUNC_RE = re.compile(r'^[၀-၉ဝ]{1,2}[။.]$')  # "၁။", "၂။"
def is_tail_garbage(s):
    """Check if line is OCR garbage at bottom: only contains :, ", \\, ', ., space."""
    ok_chars = set(':"\\\' .')
    return s and all(c in ok_chars for c in s)
HEADER_LINE_CUTOFF = 3


def is_header_line(stripped, compact):
    """Check if a line (within first N lines of a page) is a header/artifact."""
    if compact == 'မဟာဝိဟာရကျောင်းတိုက်(လှည်းကူး)':
        return True
    if compact == 'ဝတ်ရွတ်စဉ်':
        return True
    if ARTIFACT_RE.match(stripped):
        return True
    if ARTIFACT_NUM_RE.match(compact):
        return True
    if ARTIFACT_NUM_PUNC_RE.match(compact):
        return True
    return False


def clean_page_body(body, page_num):
    """Clean one page's body. Returns (cleaned_body, header_count)."""
    lines = body.split('\n')

    # PASS 1: Remove header lines (first N lines)
    remaining = []
    header_count = 0
    for idx, line in enumerate(lines):
        s = line.strip()
        c = s.replace(' ', '')
        if not s:
            remaining.append(line)
            continue
        if idx < HEADER_LINE_CUTOFF and is_header_line(s, c):
            header_count += 1
            continue
        remaining.append(line)

    # Trim leading empties
    while remaining and not remaining[0].strip():
        remaining.pop(0)

    # PASS 2: Remove watermark + tail garbage lines
    after_wm = []
    for line in remaining:
        s = line.strip()
        if s in WATERMARK_LINES:
            continue
        if is_tail_garbage(s):
            continue
        after_wm.append(line)

    # Trim trailing empties
    while after_wm and not after_wm[-1].strip():
        after_wm.pop()

    # PASS 3: Remove trailing single-char artifact lines (like "L", "+")
    cleaned = []
    skip_bottom = True
    for line in reversed(after_wm):
        if skip_bottom:
            s = line.strip()
            if ARTIFACT_RE.match(s):
                continue
            else:
                skip_bottom = False
        cleaned.insert(0, line)

    return '\n'.join(cleaned), header_count


def clean_md_file(in_path, out_path):
    with open(in_path, 'r', encoding='utf-8') as f:
        content = f.read()

    pages = content.split('## PAGE ')
    new_pages = []
    total_header = 0

    for i, page in enumerate(pages):
        if i == 0:
            new_pages.append(page)
            continue

        page_num, _, body = page.partition('\n')
        body_clean, hdr = clean_page_body(body, page_num)
        total_header += hdr
        new_pages.append(f"## PAGE {page_num}\n\n{body_clean}\n")

    result = ''.join(new_pages)

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write(result)

    removed_wm = content.count('Scanned with') - result.count('Scanned with')
    removed_maha = content.count('မဟာဝိဟာရကျောင်းတိုက်(လှည်းကူး)') - result.count('မဟာဝိဟာရကျောင်းတိုက်(လှည်းကူး)')
    removed_wat = content.count('\nဝတ်ရွတ်စဉ်\n') - result.count('\nဝတ်ရွတ်စဉ်\n')

    return {
        'watermarks': removed_wm,
        'book_title': removed_maha,
        'wat_line': removed_wat,
        'headers_removed': total_header,
        'orig_lines': len(content.split('\n')),
        'clean_lines': len(result.split('\n')),
    }


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 cleanup_mahavihara.py <input_dir> <output_dir>")
        sys.exit(1)

    input_dir = sys.argv[1]
    output_dir = sys.argv[2]

    md_files = sorted(glob.glob(os.path.join(input_dir, '*.md')))
    if not md_files:
        print(f"❌ Không tìm thấy file .md trong {input_dir}")
        sys.exit(1)

    print(f"📄 Cleanup Mahavihara: {len(md_files)} files")
    print(f"   Xóa: book title, section header, artifacts, watermark, tail garbage")
    print()

    totals = {'watermarks': 0, 'book_title': 0, 'wat_line': 0,
              'headers_removed': 0, 'orig_lines': 0, 'clean_lines': 0}

    for mf in md_files:
        out = os.path.join(output_dir, os.path.basename(mf))
        stats = clean_md_file(mf, out)

        for k in totals:
            totals[k] += stats.get(k, 0)

        tag = "✅" if stats.get('headers_removed', 0) > 0 else "▫️"
        print(f"   {tag} {os.path.basename(mf)} "
              f"(xóa {stats['headers_removed']} header, "
              f"{stats['watermarks']} watermark)")

    print()
    print(f"🎉 Done! → {os.path.abspath(output_dir)}")
    print(f"   Tổng: {len(md_files)} files cleaned")
    print(f"   Header lines xóa: {totals['headers_removed']}")
    print(f"   Watermark xóa: {totals['watermarks']}")
    print(f"   Dòng: {totals['orig_lines']} → {totals['clean_lines']}")
    print(f"   (giảm {totals['orig_lines'] - totals['clean_lines']} dòng)")


if __name__ == '__main__':
    main()