# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""ASTExtractor: language detection + dispatch to per-language extractors."""

import importlib
from pathlib import Path
from typing import Dict, Optional

from openviking.parse.parsers.code.ast.languages.base import LanguageExtractor
from openviking.parse.parsers.code.ast.skeleton import CodeSkeleton
from openviking_cli.utils import get_logger

logger = get_logger(__name__)

# File extension → internal language key
_EXT_MAP: Dict[str, str] = {
    ".py": "python",
    ".js": "javascript",
    ".jsx": "javascript",
    ".ts": "typescript",
    ".tsx": "typescript",
    ".java": "java",
    ".c": "cpp",
    ".cpp": "cpp",
    ".cc": "cpp",
    ".h": "cpp",
    ".hpp": "cpp",
    ".rs": "rust",
    ".go": "go",
    ".cs": "csharp",
    ".php": "php",
    ".lua": "lua",
}

# Language key → (module path, class name, constructor kwargs)
_EXTRACTOR_REGISTRY: Dict[str, tuple] = {
    "python": ("openviking.parse.parsers.code.ast.languages.python", "PythonExtractor", {}),
    "javascript": (
        "openviking.parse.parsers.code.ast.languages.js_ts",
        "JsTsExtractor",
        {"lang": "javascript"},
    ),
    "typescript": (
        "openviking.parse.parsers.code.ast.languages.js_ts",
        "JsTsExtractor",
        {"lang": "typescript"},
    ),
    "java": ("openviking.parse.parsers.code.ast.languages.java", "JavaExtractor", {}),
    "cpp": ("openviking.parse.parsers.code.ast.languages.cpp", "CppExtractor", {}),
    "rust": ("openviking.parse.parsers.code.ast.languages.rust", "RustExtractor", {}),
    "go": ("openviking.parse.parsers.code.ast.languages.go", "GoExtractor", {}),
    "csharp": ("openviking.parse.parsers.code.ast.languages.csharp", "CSharpExtractor", {}),
    "php": ("openviking.parse.parsers.code.ast.languages.php", "PhpExtractor", {}),
    "lua": ("openviking.parse.parsers.code.ast.languages.lua", "LuaExtractor", {}),
}


class ASTExtractor:
    """Dispatches to per-language tree-sitter extractors for supported languages.

    Unsupported languages return None, signalling the caller to fall back to LLM.
    """

    def __init__(self):
        self._cache: Dict[str, Optional[LanguageExtractor]] = {}

    def _detect_language(self, file_name: str) -> Optional[str]:
        path = Path(file_name)
        lang = _EXT_MAP.get(path.suffix.lower())
        if lang is not None:
            return lang
        # Viking resource convention: a file uploaded as <NAME.ext> is stored as a
        # directory <NAME.ext>/ containing the content under a "<NAME>.md" body
        # (e.g. foo.py/foo.md).  When the file's own suffix is unsupported (.md),
        # fall back to the parent directory's suffix to recover the language.
        if path.suffix.lower() == ".md" and path.parent.name:
            return _EXT_MAP.get(Path(path.parent.name).suffix.lower())
        return None

    def _get_extractor(self, lang: Optional[str]) -> Optional[LanguageExtractor]:
        if lang is None or lang not in _EXTRACTOR_REGISTRY:
            return None

        if lang in self._cache:
            return self._cache[lang]

        module_path, class_name, kwargs = _EXTRACTOR_REGISTRY[lang]
        try:
            mod = importlib.import_module(module_path)
            cls = getattr(mod, class_name)
            extractor = cls(**kwargs)
            self._cache[lang] = extractor
            return extractor
        except Exception as e:
            logger.warning(
                "AST extractor unavailable for language '%s', falling back to LLM: %s", lang, e
            )
            self._cache[lang] = None
            return None

    def supports(self, file_name: str) -> bool:
        """True if the file extension maps to a known language (parsing may still fail)."""
        return self._detect_language(file_name) is not None

    def extract(self, file_name: str, content: str) -> Optional[CodeSkeleton]:
        """Return raw CodeSkeleton, or None for unsupported language / parse failure.

        Consumers that need the structured object (line numbers, traversal) call
        this; consumers that want the legacy formatted text use extract_skeleton.
        """
        lang = self._detect_language(file_name)
        extractor = self._get_extractor(lang)
        if extractor is None:
            return None

        try:
            return extractor.extract(file_name, content)
        except Exception as e:
            logger.warning("AST extraction failed for '%s' (language: %s): %s", file_name, lang, e)
            return None

    def extract_skeleton(
        self, file_name: str, content: str, verbose: bool = False
    ) -> Optional[str]:
        """Extract skeleton text from source code.

        Returns None for unsupported languages or on extraction failure,
        signalling the caller to fall back to LLM.

        Args:
            verbose: If True, include full docstrings (for ast_llm / LLM input).
                     If False, only first line of each docstring (for ast / embedding).
        """
        skeleton = self.extract(file_name, content)
        if skeleton is None:
            return None
        return skeleton.to_text(verbose=verbose)


# Module-level singleton
_extractor: Optional[ASTExtractor] = None


def get_extractor() -> ASTExtractor:
    global _extractor
    if _extractor is None:
        _extractor = ASTExtractor()
    return _extractor
