#!/usr/bin/env python3
"""
Join OCR-imposed line breaks within paragraphs for Myanmar text.
Myanmar script has no inter-word spaces, so we simply concatenate.

Rules:
- Blank lines = paragraph boundary
- Lines starting with #, ---, >, -, 1., * = structural (keep as-is)
- Everything else within a paragraph = join
"""

import re
import os
import sys

# Lines that start structural elements — never join with previous
# Lines that start structural elements or table rows — never join with other lines
STRUCTURAL_PATTERNS = [
    r'^\s*#{1,6}\s',   # headers
    r'^\s*---',          # horizontal rule
    r'^\s*>\s',          # blockquote
    r'^\s*-\s',          # unordered list
    r'^\s*\d+[.)]\s',   # ordered list
    r'^\s*\*\s',        # asterisk list
    r'^\s*\|',           # table row
    r'^<!--',             # HTML comment
]
STRUCTURAL_RE = re.compile('|'.join(STRUCTURAL_PATTERNS))

# Lines that are pure horizontal rule or setext underline — never join
SEPARATOR_RE = re.compile(r'^\s*(---+|===+)\s*$')


def is_structural(line: str) -> bool:
    """Check if line starts a structural markdown element."""
    return bool(STRUCTURAL_RE.match(line))


def join_paragraph_lines(text: str) -> str:
    """
    Join OCR-broken lines within paragraphs.
    Myanmar text: no spaces needed when joining.
    Preserves markdown separators (---) to avoid setext heading bugs.
    """
    lines = text.split('\n')
    output = []
    para_buffer = []

    def flush_para():
        if para_buffer:
            joined = ''.join(para_buffer)
            output.append(joined)
            para_buffer.clear()

    for line in lines:
        stripped = line.strip()

        # Blank line = paragraph boundary
        if not stripped:
            flush_para()
            output.append('')
            continue

        # Horizontal rule (---) = never join, stand alone
        if SEPARATOR_RE.match(line):
            flush_para()
            output.append(line)
            continue

        # Structural elements = flush previous paragraph, keep line as-is
        if is_structural(line):
            flush_para()
            output.append(line)
            continue

        # Regular text line = add to paragraph buffer
        para_buffer.append(stripped)

    flush_para()

    # Remove trailing empty lines (keep at most one)
    while len(output) > 1 and output[-1] == '' and output[-2] == '':
        output.pop()

    return '\n'.join(output)


def process_file(path: str, out_path: str = None, in_place: bool = False):
    """Process a single file."""
    with open(path, 'r', encoding='utf-8') as f:
        original = f.read()

    cleaned = join_paragraph_lines(original)

    if in_place:
        with open(path, 'w', encoding='utf-8') as f:
            f.write(cleaned)
        # Stats
        orig_lines = original.count('\n')
        new_lines = cleaned.count('\n')
        print(f"  {os.path.basename(path)}: {orig_lines} → {new_lines} lines ({orig_lines - new_lines} joined)")
    elif out_path:
        with open(out_path, 'w', encoding='utf-8') as f:
            f.write(cleaned)
        orig_lines = original.count('\n')
        new_lines = cleaned.count('\n')
        print(f"  {os.path.basename(path)}: {orig_lines} → {new_lines} lines ({orig_lines - new_lines} joined)")
    else:
        return cleaned


def process_directory(dir_path: str, out_dir: str = None):
    """Process all .md files in a directory."""
    files = sorted([f for f in os.listdir(dir_path) if f.endswith('.md')])
    
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    
    for f in files:
        in_path = os.path.join(dir_path, f)
        if out_dir:
            out_path = os.path.join(out_dir, f)
            process_file(in_path, out_path)
        else:
            process_file(in_path, in_place=True)

    total_orig = 0
    total_new = 0
    for f in files:
        p = os.path.join(out_dir or dir_path, f)
        with open(p, 'r') as fh:
            total_new += fh.read().count('\n') + 1
    
    print(f"\nDone: {len(files)} files → {out_dir or dir_path}/")


if __name__ == '__main__':
    import argparse
    ap = argparse.ArgumentParser(description='Join OCR line breaks in Myanmar markdown')
    ap.add_argument('path', help='File or directory to process')
    ap.add_argument('--out', '-o', help='Output directory (default: in-place)')
    ap.add_argument('--in-place', '-i', action='store_true', help='Modify files in-place')
    args = ap.parse_args()

    if os.path.isfile(args.path):
        process_file(args.path, args.out, args.in_place)
    elif os.path.isdir(args.path):
        process_directory(args.path, args.out)
    else:
        print(f"Error: {args.path} not found")
        sys.exit(1)