from __future__ import annotations

from typing import TYPE_CHECKING

import tree_sitter_python as tspython
from tree_sitter import Language

from vibelexity.circular_deps.models import ImportInfo
from vibelexity.circular_deps.parsers.base import ImportParser
from vibelexity.parsers.shared import parse_with_language

if TYPE_CHECKING:
    from tree_sitter import Node

PY_LANGUAGE = Language(tspython.language())
_IMPORT_NODE_TYPES = {"import_statement", "import_from_statement"}
_DYNAMIC_IMPORTS = {"__import__", "importlib.import_module"}


def parse(code: str) -> Node:
    return parse_with_language(PY_LANGUAGE, code)


class PythonImportParser(ImportParser):
    def extract_imports(self, source: str, filepath: str) -> list[ImportInfo]:
        root = parse(source)
        imports: list[ImportInfo] = []

        for node in root.children:
            if node.type == "import_statement":
                imports.extend(self._process_import_statement(node))
            elif node.type == "import_from_statement":
                imports.extend(self._process_import_from_statement(node))

        return self._filter_dynamic(imports, root)

    def _process_import_statement(self, node: Node) -> list[ImportInfo]:
        imports: list[ImportInfo] = []
        has_aliased = False

        for child in node.children:
            if child.type == "aliased_import":
                has_aliased = True
                name = ""
                for nc in child.children:
                    if nc.type == "dotted_name":
                        text = nc.text.decode() if nc.text else ""
                        name = text
                    elif nc.type == "identifier":
                        text = nc.text.decode() if nc.text else ""
                        name = text
                if name:
                    imports.append(
                        ImportInfo(
                            raw_module=name,
                            resolved_path=None,
                            import_type="absolute",
                            line=node.start_point[0] + 1,
                            is_dynamic=False,
                            node_type="import_statement",
                        )
                    )
            elif not has_aliased and child.type == "dotted_name":
                text = child.text.decode() if child.text else ""
                name = text
                if name:
                    imports.append(
                        ImportInfo(
                            raw_module=name,
                            resolved_path=None,
                            import_type="absolute",
                            line=node.start_point[0] + 1,
                            is_dynamic=False,
                            node_type="import_statement",
                        )
                    )
        return imports

    def _process_import_from_statement(self, node: Node) -> list[ImportInfo]:
        imports: list[ImportInfo] = []
        module_name = ""
        is_relative = False
        found_import = False

        for child in node.children:
            if child.type == "import":
                found_import = True
                break
            if not found_import:
                if child.type == "relative_import":
                    is_relative = True
                    text = child.text.decode() if child.text else ""
                    module_name = text
                elif child.type == "dotted_name" and not module_name:
                    text = child.text.decode() if child.text else ""
                    module_name = text
            elif child.type == "module_name":
                for mc in child.children:
                    if mc.type in ("dotted_name", "relative_import"):
                        text = mc.text.decode() if mc.text else ""
                        module_name = text

        if module_name:
            imports.append(
                ImportInfo(
                    raw_module=module_name,
                    resolved_path=None,
                    import_type="relative" if is_relative else "absolute",
                    line=node.start_point[0] + 1,
                    is_dynamic=False,
                    node_type="import_from_statement",
                )
            )
        return imports

    def _filter_dynamic(
        self, imports: list[ImportInfo], root: Node
    ) -> list[ImportInfo]:
        for node in root.children:
            if node.type == "expression_statement":
                for expr in node.children:
                    if expr.type == "call":
                        func_name = self._get_function_name(expr)
                        if func_name in _DYNAMIC_IMPORTS:
                            for imp in imports:
                                if imp.line == node.start_point[0] + 1:
                                    imp.is_dynamic = True

        return [imp for imp in imports if not imp.is_dynamic]

    def _get_function_name(self, call_node: Node) -> str:
        for child in call_node.children:
            if child.type == "attribute":
                for gc in child.children:
                    if gc.type == "identifier":
                        text = gc.text.decode() if gc.text else ""
                        name = text
                        for nc in child.children:
                            text2 = nc.text.decode() if nc.text else ""
                            if nc.type == "identifier" and nc != gc:
                                parent_name = text2
                                return f"{parent_name}.{name}"
            elif child.type == "identifier":
                text = child.text.decode() if child.text else ""
                return text
        return ""
