#!/usr/bin/env python3
"""Fix <mark> tags: strip all existing, re-mark from notes (multi-format)."""
import os, re, glob

NOTES_DIR = '/home/tuan-nguyen/.openclaw/workspace/003-tang-chi-bo-giang-giai/edited-notes/gemini-flash'
EDITED_DIR = '/home/tuan-nguyen/.openclaw/workspace/003-tang-chi-bo-giang-giai/edited/gemini-flash'

def parse_notes(notes_path):
    """Parse corrections from notes file. Handles TABLE and BULLET formats.
    Returns (corrections, format_type) where format_type is 'table','bullet','narrative'."""
    with open(notes_path) as f:
        content = f.read()
    
    corrections = []
    
    # Try TABLE format first
    has_table = '|---' in content
    if has_table:
        for line in content.split('\n'):
            line = line.strip()
            if not line.startswith('|'):
                continue
            parts = [p.strip() for p in line.split('|')]
            if len(parts) < 7:
                continue
            # Skip header/separator
            if parts[1] in ('#', '---', '') or not parts[1].strip().isdigit():
                continue
            try:
                error_text = parts[4].strip(' `')
                corrected = parts[5].strip(' `')
                if error_text and corrected and error_text != corrected:
                    corrections.append({'error': error_text, 'corrected': corrected})
            except (ValueError, IndexError):
                continue
        if corrections:
            return corrections, 'table'
    
    # Try BULLET format: lines with `error` -> `corrected`
    for line in content.split('\n'):
        # Pattern: `text` -> `text` or `text` → `text`
        m = re.findall(r'`([^`]+)`\s*(?:->|→)\s*`([^`]+)`', line)
        for err, corr in m:
            err = err.strip()
            corr = corr.strip()
            if err and corr and err != corr and len(corr) >= 2:
                corrections.append({'error': err, 'corrected': corr})
    
    if corrections:
        return corrections, 'bullet'
    
    return [], 'narrative'

def strip_marks(text):
    text = re.sub(r'<mark[^>]*>', '', text)
    text = re.sub(r'</mark>', '', text)
    return text

def clean_for_search(text):
    """Remove markdown formatting for searching."""
    t = text
    t = re.sub(r'\*\*', '', t)
    t = re.sub(r'^####\s*', '', t)
    t = re.sub(r'^###\s*', '', t)
    return t.strip()

def apply_marks(content, corrections):
    """Apply <mark> tags for each correction found in content."""
    total = 0
    for corr in corrections:
        search = clean_for_search(corr['corrected'])
        if len(search) < 2:
            continue
        if search.startswith('####') or search.startswith('###'):
            continue
        
        # Only mark if the corrected text appears and is NOT already marked
        escaped = re.escape(search)
        # Find first occurrence and mark it
        pattern = r'(?<!<mark>)(' + escaped + r')(?!</mark>)'
        new_content, n = re.subn(pattern, r'<mark>\1</mark>', content, count=1)
        if n > 0:
            content = new_content
            total += n
    return content, total

def fix_batch(batch_name):
    """Fix one batch."""
    edited_path = os.path.join(EDITED_DIR, f'{batch_name}.md')
    notes_path = os.path.join(NOTES_DIR, f'{batch_name}-notes.md')
    
    if not os.path.exists(edited_path):
        return False, "no edited file"
    if not os.path.exists(notes_path):
        return False, "no notes file"
    
    with open(edited_path) as f:
        content = f.read()
    
    corrections, fmt = parse_notes(notes_path)
    
    old_marks = len(re.findall(r'<mark', content))
    content = strip_marks(content)
    
    if fmt == 'narrative':
        # Strip only - no specific corrections to re-mark
        with open(edited_path, 'w') as f:
            f.write(content)
        return True, f"stripped {old_marks} marks (narrative notes, no re-mark)"
    
    # Re-apply marks
    content, marked = apply_marks(content, corrections)
    
    with open(edited_path, 'w') as f:
        f.write(content)
    
    return True, f"{fmt}: {len(corrections)} corr → {marked} marked (was {old_marks})"

# Main
batches = [
    'output-192-to-194', 'output-237-to-239', 'output-276-to-278',  # excessive
    'output-207-to-209',  # heavy
    'output-198-to-200', 'output-201-to-203', 'output-219-to-221', 'output-228-to-230',  # no marks
    'output-243-to-245', 'output-249-to-251', 'output-252-to-254', 'output-267-to-269',
    'output-282-to-284', 'output-294-to-296',
    'output-195-to-197', 'output-234-to-236', 'output-240-to-242', 'output-246-to-248',  # few marks
    'output-258-to-260', 'output-261-to-263', 'output-279-to-281', 'output-291-to-293'
]

print(f"Fixing {len(batches)} batches...\n")
ok = fail = 0
for batch in batches:
    success, msg = fix_batch(batch)
    print(f"  {'✅' if success else '❌'} {batch}: {msg}")
    if success: ok += 1
    else: fail += 1

print(f"\nDone: {ok} fixed, {fail} failed")

# Final verify
print("\n=== Final marks ===")
for batch in sorted(batches):
    path = os.path.join(EDITED_DIR, f'{batch}.md')
    if os.path.exists(path):
        with open(path) as f:
            marks = len(re.findall(r'<mark[^>]*>', f.read()))
        print(f"  {batch}: {marks} marks")