#!/usr/bin/env python3
"""
Setup script to install tree-sitter language grammars via pip.

Usage:
    python syntax/setup_grammars.py           # Install bundled languages (python, c, cpp)
    python syntax/setup_grammars.py rust go   # Install specific languages
    python syntax/setup_grammars.py --all     # Install all known languages
    python syntax/setup_grammars.py --list    # List available languages
"""

import sys
import subprocess

# Mapping from language name to pip package
LANGUAGE_PACKAGES = {
    'python': 'tree-sitter-python',
    'c': 'tree-sitter-c',
    'cpp': 'tree-sitter-cpp',
    'javascript': 'tree-sitter-javascript',
    'typescript': 'tree-sitter-typescript',
    'rust': 'tree-sitter-rust',
    'go': 'tree-sitter-go',
    'java': 'tree-sitter-java',
    'json': 'tree-sitter-json',
    'bash': 'tree-sitter-bash',
    'markdown': 'tree-sitter-markdown',
}

# Bundled languages (installed by default)
BUNDLED = ['python', 'c', 'cpp']


def check_installed(lang: str) -> bool:
    """Check if a language is already installed."""
    module_name = LANGUAGE_PACKAGES[lang].replace('-', '_')
    try:
        __import__(module_name)
        return True
    except ImportError:
        return False


def install_language(lang: str) -> bool:
    """Install a language via pip."""
    if lang not in LANGUAGE_PACKAGES:
        print(f"Unknown language: {lang}")
        return False

    package = LANGUAGE_PACKAGES[lang]
    print(f"[INSTALL] {lang} ({package})")

    try:
        result = subprocess.run(
            [sys.executable, '-m', 'pip', 'install', package],
            capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"  ERROR: {result.stderr.strip()}")
            return False
        print(f"  OK")
        return True
    except Exception as e:
        print(f"  ERROR: {e}")
        return False


def main():
    if len(sys.argv) > 1:
        if sys.argv[1] == '--list':
            print("Available languages:")
            for lang in sorted(LANGUAGE_PACKAGES.keys()):
                status = "[installed]" if check_installed(lang) else ""
                bundled = "(bundled)" if lang in BUNDLED else ""
                print(f"  {lang:15} {bundled:10} {status}")
            return
        elif sys.argv[1] == '--all':
            languages = list(LANGUAGE_PACKAGES.keys())
        else:
            languages = sys.argv[1:]
    else:
        # Default: bundled languages
        languages = BUNDLED

    # Check for tree-sitter
    try:
        import tree_sitter
    except ImportError:
        print("ERROR: tree-sitter not installed.")
        print("Run: pip install tree-sitter")
        sys.exit(1)

    print(f"Installing languages: {', '.join(languages)}")
    print()

    installed = []
    failed = []

    for lang in languages:
        if lang not in LANGUAGE_PACKAGES:
            print(f"Unknown language: {lang}")
            failed.append(lang)
            continue

        if check_installed(lang):
            print(f"[SKIP] {lang} - already installed")
            installed.append(lang)
            continue

        if install_language(lang):
            installed.append(lang)
        else:
            failed.append(lang)

    print()
    print("Done!")
    print(f"Installed: {installed}")
    if failed:
        print(f"Failed: {failed}")


if __name__ == '__main__':
    main()
