#!/usr/bin/env python3
"""
apply-marks.py — Đánh dấu <mark> từ notes vào edited files

Dựa vào cột "Đã sửa thành" trong edited-notes → tìm trong file edited → bọc <mark>.
"""

import re, os, sys, glob, argparse
from collections import defaultdict

PROJ = os.path.dirname(os.path.abspath(__file__))
EDITED = f"{PROJ}/edited/gemini-flash"
NOTES = f"{PROJ}/edited-notes/gemini-flash"
BACKUP = f"{PROJ}/_backup/marks"


def parse_notes(note_path: str) -> list[str]:
    """Extract corrected terms from the notes markdown table."""
    if not os.path.exists(note_path):
        print(f"  ⚠️  Notes not found: {note_path}")
        return []

    with open(note_path) as f:
        content = f.read()

    lines = content.strip().split('\n')
    terms = []
    corrected_idx = None
    header_found = False

    # ---- Method 1: Standard pipe table ----
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith('|') and ('Đã sửa thành' in stripped or 'Corrected To' in stripped):
            header_found = True
            # Found header — determine column index
            cells = [c.strip() for c in stripped.split('|')]
            for j, c in enumerate(cells):
                if 'Đã sửa thành' in c or 'Corrected To' in c:
                    corrected_idx = j
                    break

            # Parse data rows after the separator line
            data_start = i + 2  # Skip header + separator
            for row_line in lines[data_start:]:
                row_line = row_line.strip()
                if not row_line.startswith('|'):
                    break  # End of table
                if re.match(r'^\|[- ]+\|', row_line):
                    continue  # Another separator
                cells = [c.strip() for c in row_line.split('|')]
                if corrected_idx < len(cells):
                    corrected = cells[corrected_idx]
                    # Clean up: strip bold, backticks, parenthetical notes, HTML
                    corrected = corrected.replace('**', '').replace('`', '').replace('<br>', '')
                    corrected = re.sub(r'\([^)]*\)', '', corrected)
                    corrected = corrected.strip()
                    if corrected:
                        terms.append(corrected)
            return terms  # Done parsing table

    # ---- Method 2: Arrow format (e.g., `term -> corrected_term`) ----
    for line in lines:
        stripped = line.strip()
        if not stripped or stripped.startswith('#'):
            continue
        # Match patterns like: `old -> new` or `old → new` or `"old" -> "new"`
        # Try backtick-quoted pattern: `old` -> `new` (desc)
        m = re.match(r'^(?:[-*]\s*)?`([^`]+)`\s*(?:->|→)\s*`([^`]+)`', stripped)
        if not m:
            # Try unquoted: old -> new  or  old → new
            m = re.match(r'^(?:[-*]\s*)?(?:`|\")?([^`"\s][^`~"→-]*?)(?:`|\")?\s*(?:->|→|=>)\s*(?:`|\**)?([^`\s][^`";→-]*?)(?:`|\*+)?', stripped)
        if m:
            old, new = m.group(1).strip(), m.group(2).strip()
            # Skip if old == new or new is just formatting
            if len(new) >= 2 and old != new:
                terms.append(new)

    return terms


def apply_marks(text: str, terms: list[str]) -> tuple[str, int]:
    """
    Wrap all occurrences of each term in <mark> tags.
    Returns (marked_text, total_marks_applied).
    - Sorts terms by length (desc) so longer terms are processed first
    - Handles overlapping by keeping the first/longer match
    - Avoids double-wrapping (text already inside <mark> is skipped)
    """
    # Filter and sort: longest first
    unique_terms = sorted(
        {t for t in terms if len(t) >= 2},
        key=lambda t: -len(t)
    )

    # Phase 1: collect all match positions
    all_hits = []  # (start, end, term)
    for term in unique_terms:
        esc = re.escape(term)
        for match in re.finditer(esc, text):
            all_hits.append((match.start(), match.end(), term))

    # Phase 2: deduplicate overlapping matches — prefer longer
    all_hits.sort(key=lambda h: (h[0], -(h[1] - h[0])))  # start asc, length desc
    keep = []
    for start, end, term in all_hits:
        overlap = False
        for ks, ke, _ in keep:
            if start < ke and end > ks:  # overlap
                overlap = True
                break

        # Also skip if inside existing <mark> tag
        before = text[max(0, start - 10):start]
        after = text[end:end + 10]
        if '<mark>' in before or '</mark>' in after:
            # Check more carefully
            chunk_before = text[max(0, start - 200):start]
            open_marks = chunk_before.count('<mark>') - chunk_before.count('</mark>')
            if open_marks > 0:
                overlap = True

        if not overlap:
            keep.append((start, end, term))

    # Phase 3: apply from right-to-left
    keep.sort(key=lambda h: -h[0])
    result = list(text)
    for start, end, term in keep:
        result[start:end] = list(f'<mark>{term}</mark>')

    return ''.join(result), len(keep)


def process_one(batch_name: str, dry_run: bool = False) -> dict:
    """Process a single batch. Returns stats dict."""
    edit_path = f"{EDITED}/{batch_name}.md"
    note_path = f"{NOTES}/{batch_name}-notes.md"

    if not os.path.exists(edit_path):
        return {"name": batch_name, "error": "edited file not found"}
    if not os.path.exists(note_path):
        return {"name": batch_name, "error": "notes file not found"}

    with open(edit_path) as f:
        original = f.read()

    terms = parse_notes(note_path)
    if not terms:
        return {"name": batch_name, "error": "no terms found in notes"}

    # Split terms that have spaces into individual tokens
    all_tokens = []
    for t in terms:
        for token in t.split():
            token = token.strip()
            if token and len(token) >= 2:
                all_tokens.append(token)

    marked, count = apply_marks(original, all_tokens)

    if dry_run:
        return {
            "name": batch_name,
            "dry_run": True,
            "terms_found": len(terms),
            "tokens_total": len(all_tokens),
            "marks_applied": count,
            "original_size": len(original),
            "marked_size": len(marked),
            "increase": len(marked) - len(original),
        }

    # Backup original
    os.makedirs(BACKUP, exist_ok=True)
    import time
    ts = time.strftime("%Y%m%d_%H%M%S")
    backup_path = f"{BACKUP}/{batch_name}.md.bak.{ts}"
    with open(backup_path, 'w') as f:
        f.write(original)

    # Write marked version
    with open(edit_path, 'w') as f:
        f.write(marked)

    return {
        "name": batch_name,
        "terms_found": len(terms),
        "tokens_total": len(all_tokens),
        "marks_applied": count,
        "original_size": len(original),
        "marked_size": len(marked),
        "increase": len(marked) - len(original),
        "backup": backup_path,
    }


def process_all(dry_run: bool = False):
    """Process all batches that have both edited and note files."""
    edited_files = {os.path.basename(f).replace('.md', '')
                    for f in glob.glob(f"{EDITED}/*.md") if not f.endswith('-notes.md')}
    note_files = {os.path.basename(f).replace('-notes.md', '')
                  for f in glob.glob(f"{NOTES}/*-notes.md")}

    common = sorted(edited_files & note_files)

    results = []
    for name in common:
        r = process_one(name, dry_run=dry_run)
        results.append(r)
        status = "DRY" if dry_run else "OK"
        if "error" in r:
            print(f"  ❌ {name}: {r['error']}")
        else:
            print(f"  {status} {name}: {r['marks_applied']} marks ({r['tokens_total']} tokens from {r['terms_found']} terms)")

    # Summary
    total_marks = sum(r.get('marks_applied', 0) for r in results)
    errors = [r for r in results if 'error' in r]
    print(f"\n{'DRY RUN — ' if dry_run else ''}Summary: {len(results)} files, {total_marks} marks, {len(errors)} errors")
    return results


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch', help='Process single batch (e.g. output-6-to-10)')
    parser.add_argument('--dry-run', action='store_true', help='Preview without modifying files')
    parser.add_argument('--all', action='store_true', help='Process all batches')
    parser.add_argument('--list-file', help='File with batch names (one per line)')
    args = parser.parse_args()

    if args.batch:
        r = process_one(args.batch, dry_run=args.dry_run)
        if 'error' in r:
            print(f"❌ {r['name']}: {r['error']}")
        else:
            print(f"{'DRY RUN — ' if args.dry_run else ''}Batch: {r['name']}")
            print(f"  Terms from notes:  {r['terms_found']}")
            print(f"  Tokens to find:    {r['tokens_total']}")
            print(f"  Marks applied:     {r['marks_applied']}")
            print(f"  Size: {r['original_size']:,} → {r['marked_size']:,} (+{r['increase']:,} chars)")
    elif args.list_file:
        with open(args.list_file) as f:
            names = [line.strip() for line in f if line.strip()]
        results = []
        for name in names:
            r = process_one(name, dry_run=args.dry_run)
            results.append(r)
            if 'error' in r:
                print(f"  ❌ {name}: {r['error']}")
            else:
                print(f"  {'DRY' if args.dry_run else 'OK'} {name}: {r['marks_applied']} marks")
        total = sum(r.get('marks_applied', 0) for r in results)
        errs = [r for r in results if 'error' in r]
        print(f"\n{'DRY RUN — ' if args.dry_run else ''}Summary: {len(results)} files, {total} marks, {len(errs)} errors")
    elif args.all:
        process_all(dry_run=args.dry_run)
    else:
        print("Usage: python3 apply-marks.py --batch output-6-to-10 [--dry-run]")
        print("       python3 apply-marks.py --all [--dry-run]")