from typing import Optional, Dict
import tree_sitter
import tree_sitter_python
import tree_sitter_javascript
import tree_sitter_typescript
import tree_sitter_java
import tree_sitter_php
import tree_sitter_html
import tree_sitter_json
import tree_sitter_bash
from tree_sitter import Language, Parser

SUPPORTED_LANGUAGES = {
    ".py": "python",
    ".js": "javascript",
    ".ts": "typescript",
    ".tsx": "typescript",
    ".php": "php",
    ".html": "html",
    ".json": "json",
    ".sh": "bash",
    ".bash": "bash",
}

# Cache for loaded languages
_LANG_CACHE: Dict[str, Language] = {}

def get_lang_for_extension(ext: str) -> Optional[str]:
    return SUPPORTED_LANGUAGES.get(ext)

def load_language(language_name: str) -> Language:
    if language_name in _LANG_CACHE:
        return _LANG_CACHE[language_name]
    
    if language_name == "python":
        lang = Language(tree_sitter_python.language())
    elif language_name == "javascript":
        lang = Language(tree_sitter_javascript.language())
    elif language_name == "typescript":
        lang = Language(tree_sitter_typescript.language())
    elif language_name == "php":
        lang = Language(tree_sitter_php.language_php())
    elif language_name == "html":
        lang = Language(tree_sitter_html.language())
    elif language_name == "json":
        lang = Language(tree_sitter_json.language())
    elif language_name == "bash":
        lang = Language(tree_sitter_bash.language())
    else:
        raise ValueError(f"Language {language_name} is not supported.")
    
    _LANG_CACHE[language_name] = lang
    return lang

def parse_code(code: bytes, language_name: str) -> tree_sitter.Tree:
    lang = load_language(language_name)
    parser = Parser(lang)
    return parser.parse(code)
