"""Idempotent patcher for leann's CODE_EXTENSIONS in chunking_utils.py.

Adds file extension mappings for languages supported by astchunk-extended
but not present in upstream leann. Safe to run multiple times.

Usage:
    python -m astchunk.patch_leann
    python -c "from astchunk.patch_leann import apply; apply()"
"""

import importlib.util
import re
import sys
from pathlib import Path

EXTRA_EXTENSIONS = {
    '".c"': '"c"',
    '".h"': '"c"',
    '".cpp"': '"cpp"',
    '".cc"': '"cpp"',
    '".cxx"': '"cpp"',
    '".hpp"': '"cpp"',
    '".hh"': '"cpp"',
    '".go"': '"go"',
    '".rs"': '"rust"',
    '".rb"': '"ruby"',
    '".sh"': '"bash"',
    '".bash"': '"bash"',
    '".zsh"': '"bash"',
    '".kt"': '"kotlin"',
    '".kts"': '"kotlin"',
    '".scala"': '"scala"',
    '".sc"': '"scala"',
    '".lua"': '"lua"',
    '".php"': '"php"',
    '".zig"': '"zig"',
}


def _find_chunking_utils() -> Path | None:
    """Locate leann's chunking_utils.py via importlib or known paths."""
    spec = importlib.util.find_spec("leann")
    if spec and spec.submodule_search_locations:
        for loc in spec.submodule_search_locations:
            candidate = Path(loc) / "chunking_utils.py"
            if candidate.exists():
                return candidate

    # Fallback: search in uv tools directory
    uv_tools = Path.home() / ".local" / "share" / "uv" / "tools" / "leann-core"
    if uv_tools.exists():
        matches = list(uv_tools.rglob("leann/chunking_utils.py"))
        if matches:
            return matches[0]

    return None


def _is_already_patched(content: str) -> bool:
    """Check if CODE_EXTENSIONS already contains our additions."""
    return '".c"' in content and '".go"' in content and '".rs"' in content


def _build_patch_lines() -> str:
    """Generate the extra extension lines to insert."""
    lines = []
    for ext, lang in EXTRA_EXTENSIONS.items():
        lines.append(f"    {ext}: {lang},")
    return "\n".join(lines)


def apply(verbose: bool = True) -> bool:
    """Apply the CODE_EXTENSIONS patch to leann's chunking_utils.py.

    Returns:
        True if patch was applied, False if already patched or file not found.
    """
    target = _find_chunking_utils()
    if target is None:
        if verbose:
            print("leann chunking_utils.py not found, skipping patch")
        return False

    content = target.read_text()

    if _is_already_patched(content):
        if verbose:
            print(f"Already patched: {target}")
        return False

    # Find the closing brace of CODE_EXTENSIONS dict
    # Pattern: last entry before the closing }
    pattern = r'(CODE_EXTENSIONS\s*=\s*\{[^}]*?)(\.jsx["\']:\s*["\']typescript["\'],?\s*\n)(\})'
    match = re.search(pattern, content)

    if not match:
        # Fallback: find any closing } after CODE_EXTENSIONS
        pattern2 = r'(CODE_EXTENSIONS\s*=\s*\{[^}]*?)(\})'
        match = re.search(pattern2, content, re.DOTALL)
        if not match:
            if verbose:
                print(f"Could not find CODE_EXTENSIONS dict in {target}")
            return False
        new_content = content[:match.end(1)] + _build_patch_lines() + "\n" + content[match.start(2):]
    else:
        insert_point = match.end(2)
        new_content = content[:insert_point] + _build_patch_lines() + "\n" + content[insert_point:]

    target.write_text(new_content)
    if verbose:
        print(f"Patched {target}: added {len(EXTRA_EXTENSIONS)} extension mappings")
    return True


if __name__ == "__main__":
    success = apply(verbose=True)
    sys.exit(0 if success else 1)
