"""
Tree-sitter based parser — language-agnostic AST extraction.

Dynamically loads grammar packages for each supported language and walks
the concrete syntax tree to extract classes, functions, imports, etc.
"""

from __future__ import annotations

import logging
import re
import threading
from pathlib import Path
from typing import Any, Optional

from hld_generator.config import TREE_SITTER_GRAMMARS
from hld_generator.parsers.base import (
    EntityType,
    ImportInfo,
    ParsedEntity,
    ParsedFile,
    extract_summary,
)

logger = logging.getLogger(__name__)

# ── Tree-sitter query patterns per language ────────────────────────────────────
# Each value is a dict mapping EntityType keys to S-expression queries.

_QUERIES: dict[str, dict[str, str]] = {
    "python": {
        "class": "(class_definition name: (identifier) @name) @class",
        "function": "(function_definition name: (identifier) @name) @func",
        "import": "[((import_statement) @imp) ((import_from_statement) @imp)]",
        "decorator": "(decorator) @dec",
    },
    "javascript": {
        "class": "(class_declaration name: (identifier) @name) @class",
        "function": "[(function_declaration name: (identifier) @name) @func (method_definition name: (property_identifier) @name) @func]",
        "import": "(import_statement) @imp",
        "export": "[(export_statement) @exp (export_default_declaration) @exp]",
    },
    "typescript": {
        "class": "(class_declaration name: (type_identifier) @name) @class",
        "function": "[(function_declaration name: (identifier) @name) @func (method_definition name: (property_identifier) @name) @func]",
        "import": "(import_statement) @imp",
        "export": "[(export_statement) @exp (export_default_declaration) @exp]",
        "interface": "(interface_declaration name: (type_identifier) @name) @iface",
    },
    "java": {
        "class": "[(class_declaration name: (identifier) @name) @class (interface_declaration name: (identifier) @name) @class]",
        "function": "(method_declaration name: (identifier) @name) @func",
        "import": "(import_declaration) @imp",
    },
    "go": {
        "class": "(type_declaration (type_spec name: (type_identifier) @name)) @class",
        "function": "[(function_declaration name: (identifier) @name) @func (method_declaration name: (field_identifier) @name) @func]",
        "import": "(import_declaration) @imp",
    },
    "rust": {
        "class": "[(struct_item name: (type_identifier) @name) @class (enum_item name: (type_identifier) @name) @class (trait_item name: (type_identifier) @name) @class]",
        "function": "[(function_item name: (identifier) @name) @func (impl_item) @impl]",
        "import": "(use_declaration) @imp",
    },
    "c": {
        "class": "(struct_specifier name: (type_identifier) @name) @class",
        "function": "(function_definition declarator: (function_declarator declarator: (identifier) @name)) @func",
        "import": "(preproc_include) @imp",
    },
    "cpp": {
        "class": "[(class_specifier name: (type_identifier) @name) @class (struct_specifier name: (type_identifier) @name) @class]",
        "function": "(function_definition declarator: (function_declarator declarator: [(identifier) @name (qualified_identifier) @name])) @func",
        "import": "[(preproc_include) @imp (using_declaration) @imp]",
    },
    "ruby": {
        "class": "[(class name: [(constant) @name (scope_resolution) @name]) @class (module name: (constant) @name) @class]",
        "function": "(method name: (identifier) @name) @func",
    },
}


class TreeSitterParser:
    """Parse source files using tree-sitter grammars."""

    def __init__(self) -> None:
        self._parsers: dict[str, Any] = {}
        self._languages: dict[str, Any] = {}
        self._lock = threading.Lock()
        self._regex_parser: Any = None  # cached RegexParser for endpoint extraction

    # ── public ─────────────────────────────────────────────────────────────

    def supports(self, language: str) -> bool:
        """Check if we can load a grammar for *language*."""
        return language in TREE_SITTER_GRAMMARS

    def parse_file(self, file_path: Path, language: str) -> ParsedFile:
        """Parse a single source file and return structured data."""
        result = ParsedFile(file_path=str(file_path), language=language)

        try:
            source = file_path.read_bytes()
            result.line_count = source.count(b"\n") + 1
        except OSError as exc:
            result.errors.append(f"Cannot read file: {exc}")
            return result

        # Extract top-of-file docstring / comment block
        result.raw_summary = extract_summary(source.decode("utf-8", errors="replace"))

        parser = self._get_parser(language)
        if parser is None:
            result.errors.append(f"No tree-sitter grammar for {language}")
            return result

        try:
            tree = parser.parse(source)
        except Exception as exc:
            result.errors.append(f"Parse error: {exc}")
            return result

        source_text = source.decode("utf-8", errors="replace")

        # Run queries
        lang_queries = _QUERIES.get(language, {})
        with self._lock:
            ts_language = self._languages[language]

        for query_kind, query_str in lang_queries.items():
            try:
                self._run_query(
                    result, tree, ts_language, source, source_text, query_kind, query_str,
                    language=language,
                )
            except Exception as exc:
                logger.debug("Query %s failed for %s: %s", query_kind, language, exc)

        # Tree-sitter queries don't cover endpoint detection — use regex patterns
        if self._regex_parser is None:
            from hld_generator.parsers.regex_parser import RegexParser
            self._regex_parser = RegexParser()
        self._regex_parser._extract_endpoints(result, source_text)
        # Detect Python `if __name__ == "__main__":` guard for entry point detection
        self._regex_parser._detect_main_guard(result, source_text)

        return result

    # ── private ────────────────────────────────────────────────────────────

    def _get_parser(self, language: str) -> Any | None:
        with self._lock:
            if language in self._parsers:
                return self._parsers[language]

        grammar_pkg = TREE_SITTER_GRAMMARS.get(language)
        if not grammar_pkg:
            return None

        try:
            import tree_sitter as ts

            mod = __import__(grammar_pkg)
            lang = ts.Language(mod.language())
            parser = ts.Parser(lang)
            with self._lock:
                self._parsers[language] = parser
                self._languages[language] = lang
            return parser
        except Exception as exc:
            logger.warning("Cannot load tree-sitter grammar for %s: %s", language, exc)
            with self._lock:
                self._parsers[language] = None
            return None

    @staticmethod
    def _node_text(source_bytes: bytes, node: Any) -> str:
        """Extract text from a tree-sitter node using byte offsets on raw bytes."""
        return source_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace")

    def _run_query(
        self,
        result: ParsedFile,
        tree: Any,
        ts_language: Any,
        source_bytes: bytes,
        source_text: str,
        kind: str,
        query_str: str,
        language: str = "",
    ) -> None:
        try:
            import tree_sitter as ts
            # tree-sitter 0.24+: Language.query() is deprecated; use Query() constructor.
            try:
                query = ts.Query(ts_language, query_str)
            except (TypeError, AttributeError):
                query = ts_language.query(query_str)
        except Exception:
            return

        captures = query.captures(tree.root_node)

        # tree-sitter 0.22+ returns dict[str, list[Node]]
        if isinstance(captures, dict):
            nodes = []
            for key in captures:
                for node in captures[key]:
                    nodes.append((node, key))
        else:
            nodes = captures

        for node, capture_name in nodes:
            if kind == "import":
                if capture_name == "imp":
                    self._extract_import(result, source_bytes, node, language)
            elif kind == "export":
                if capture_name == "exp":
                    text = self._node_text(source_bytes, node)
                    result.exports.append(text.split("\n")[0][:120])
            elif kind == "decorator":
                continue  # decorators handled with their functions
            elif kind in ("class", "interface"):
                if capture_name == "name":
                    etype = EntityType.INTERFACE if kind == "interface" else EntityType.CLASS
                    entity = ParsedEntity(
                        name=self._node_text(source_bytes, node),
                        entity_type=etype,
                        line_start=node.start_point[0] + 1,
                        line_end=node.end_point[0] + 1,
                    )
                    result.entities.append(entity)
            elif kind == "function":
                if capture_name == "name":
                    parent_class = self._find_parent_class(node, source_bytes)
                    etype = EntityType.METHOD if parent_class else EntityType.FUNCTION
                    entity = ParsedEntity(
                        name=self._node_text(source_bytes, node),
                        entity_type=etype,
                        line_start=node.start_point[0] + 1,
                        line_end=node.end_point[0] + 1,
                        parent=parent_class,
                    )
                    result.entities.append(entity)

    def _extract_import(self, result: ParsedFile, source_bytes: bytes, node: Any, language: str = "") -> None:
        text = self._node_text(source_bytes, node).strip()

        # Go import blocks: import (\n\t"pkg1"\n\t"pkg2"\n)
        # Only trigger for Go files to avoid false positives with Scala/Dart.
        if language == "go" and text.startswith("import") and "(" in text:
            self._extract_go_import_block(result, text)
            return

        imp = ImportInfo(module=text, names=[])

        # Simple heuristic extraction
        if "from " in text and " import " in text:
            parts = text.split("from ", 1)[1].split(" import ", 1)
            imp.module = parts[0].strip().strip("'\"")
            if len(parts) > 1:
                imp.names = [n.strip().split(" as ")[0].strip() for n in parts[1].split(",")]
            imp.is_relative = imp.module.startswith(".")
        elif text.startswith("import "):
            # Single import: import "fmt" (Go) or import os (Python/Java)
            raw = text.replace("import ", "").strip().split(" as ")[0].split(";")[0].strip()
            imp.module = raw.strip('"')
        elif text.startswith("#include"):
            imp.module = text.replace("#include", "").strip().strip('<>"')
        elif text.startswith("use "):
            imp.module = text.replace("use ", "").strip().rstrip(";")

        result.imports.append(imp)

    @staticmethod
    def _extract_go_import_block(result: ParsedFile, text: str) -> None:
        """Parse a Go import block into individual ImportInfo entries."""
        for match in re.finditer(r'"([^"]+)"', text):
            module = match.group(1)
            result.imports.append(ImportInfo(module=module, names=[]))

    def _find_parent_class(self, node: Any, source_bytes: bytes) -> Optional[str]:
        """Walk up the AST to find an enclosing class/struct/impl."""
        current = node.parent
        class_types = {
            "class_definition", "class_declaration", "class_specifier",
            "struct_specifier", "impl_item", "class", "module",
        }
        while current:
            if current.type in class_types:
                for child in current.children:
                    if child.type in ("identifier", "type_identifier", "constant", "name"):
                        return self._node_text(source_bytes, child)
            current = current.parent
        return None

