import logging
import sys
from collections.abc import Sequence
from pathlib import Path

import libcst as cst
import libcst.matchers as m
from libcst.metadata import PositionProvider

from py_agent.models.annotation_type import AnnotationType
from py_agent.models.parameter import RawParentTypeMetadata, RawTypeMetadata
from py_agent.models.parameter_type_kind import ParameterTypeKind
from py_agent.utils.constants import PRIMITIVE_TYPES

logger = logging.getLogger(__name__)


class typeVisitor(m.MatcherDecoratableVisitor):
    METADATA_DEPENDENCIES = (PositionProvider,)

    def __init__(
        self,
        absolute_file_path: str = "",
        project_path: str = "",
        type_names: list[str] | None = None,
    ):
        super().__init__()
        self.absolute_file_path: Path = Path(absolute_file_path)
        self.project_path: Path = Path(project_path)
        self.type_names: list[str] = type_names or []
        # Lazy import to avoid circular dependency
        from py_agent.services.import_jedi_extractor import ImportJediExtractor

        self.jedi = ImportJediExtractor(absolute_file_path=str(self.absolute_file_path))

        self.visited_classes: set[int] = set()
        self.type_mapping: dict[str, RawTypeMetadata] = {}

    def get_type_mapping(self) -> dict[str, RawTypeMetadata]:
        return self.type_mapping

    def _is_core(self, path: Path) -> bool:
        return path.is_relative_to(self.jedi.get_env_path()) or path.is_relative_to(
            Path(sys.base_prefix)
        )

    def _is_current_file_core(self) -> bool:
        return self._is_core(self.absolute_file_path)

    def _safe_name_to_module_name(self, name: str) -> str:
        metadata = self.jedi.get_jedi_metadata(name)
        if not metadata:
            return ""
        return metadata.module_name

    def _recursive_unravel(self, node: cst.CSTNode) -> set[AnnotationType]:
        if isinstance(node, cst.Name):
            name = node.value
            pos = self.get_metadata(PositionProvider, node, None)
            kind = ParameterTypeKind.PRIMITIVE if name in PRIMITIVE_TYPES else ParameterTypeKind.ANY
            line = pos.end.line if pos else -1
            col = pos.end.column if pos else -1
            return {AnnotationType(name, kind, namespace=None, full_name=name, line=line, col=col)}

        if isinstance(node, cst.Attribute):
            full_name = cst.Module([]).code_for_node(node)
            namespace = cst.Module([]).code_for_node(node.value)
            pos = self.get_metadata(PositionProvider, node.attr, None)
            name = node.attr.value
            kind = ParameterTypeKind.PRIMITIVE if name in PRIMITIVE_TYPES else ParameterTypeKind.ANY
            line = pos.end.line if pos else -1
            col = pos.end.column if pos else -1
            return {
                AnnotationType(
                    name, kind, namespace=namespace, full_name=full_name, line=line, col=col
                )
            }

        if isinstance(node, cst.Annotation):
            return self._recursive_unravel(node.annotation)

        if isinstance(node, cst.Index):
            return self._recursive_unravel(node.value)

        if isinstance(node, cst.BinaryOperation):
            return self._recursive_unravel(node.left) | self._recursive_unravel(node.right)

        if isinstance(node, cst.Subscript):
            result: set[AnnotationType] = set()
            for sli in node.slice:
                result.update(self._recursive_unravel(sli))
            return result

        if isinstance(node, cst.SubscriptElement):
            return self._recursive_unravel(node.slice)

        return set()

    @staticmethod
    def _get_annassign_fields(class_node: cst.ClassDef) -> list[cst.AnnAssign]:
        annassigns = []
        for stmt in class_node.body.body:
            if isinstance(stmt, cst.AnnAssign):
                annassigns.append(stmt)
            elif isinstance(stmt, cst.SimpleStatementLine):
                for expr in stmt.body:
                    if isinstance(expr, cst.AnnAssign):
                        annassigns.append(expr)
        return annassigns

    def _pos_to_raw_parent(self, line: int, col: int) -> RawParentTypeMetadata | None:
        jedi_metadata = self.jedi.get_by_position(line, col)
        if not jedi_metadata:
            return None
        if jedi_metadata.is_core:
            return None
        if not jedi_metadata.file_path:
            return None
        return RawParentTypeMetadata(
            name=jedi_metadata.name,
            import_from=jedi_metadata.file_path,
            is_core_module=jedi_metadata.is_core,
            project_path=self.project_path,
        )

    def _references_to_parents(self, children: list[cst.AnnAssign]) -> set[RawParentTypeMetadata]:
        result: set[RawParentTypeMetadata] = set()
        for child in children:
            types = self._recursive_unravel(child.annotation)
            for x in types:
                if x.line is None or x.col is None:
                    continue
                raw_parent = self._pos_to_raw_parent(x.line, x.col)
                if raw_parent:
                    result.add(raw_parent)
        return result

    def _bases_to_parents(self, bases: Sequence[cst.Arg]) -> set[RawParentTypeMetadata]:
        result: set[RawParentTypeMetadata] = set()
        for base in bases:
            pos = self.get_metadata(PositionProvider, base.value, None)
            if pos:
                raw_parent = self._pos_to_raw_parent(pos.end.line, pos.end.column)
                if raw_parent:
                    result.add(raw_parent)
        return result

    @m.visit(
        m.ClassDef(
            decorators=m.MatchIfTrue(
                lambda decos: any(
                    m.matches(
                        deco.decorator,
                        m.Name("dataclass")
                        | m.Attribute(attr=m.Name("dataclass"))
                        | m.Call(func=m.Name("dataclass") | m.Attribute(attr=m.Name("dataclass"))),
                    )
                    for deco in decos
                )
            )
        )
    )
    def visit_dataclass(self, node: cst.ClassDef) -> None:
        if self.type_names and node.name.value not in self.type_names:
            return
        parents = self._bases_to_parents(node.bases)
        children = self._references_to_parents(typeVisitor._get_annassign_fields(node))
        dataclass_type = RawTypeMetadata(
            name=node.name.value,
            kind=ParameterTypeKind.CLASS,
            import_from=self.absolute_file_path,
            module_name=self._safe_name_to_module_name(node.name.value),
            project_path=self.project_path,
            is_leaf_DTO=True,
            is_core_module=self._is_current_file_core(),
            parents=frozenset(parents),
            attributes=frozenset(children),
            code=cst.Module([]).code_for_node(node),
        )
        self.type_mapping[dataclass_type.name] = dataclass_type
        self.visited_classes.add(id(node))

    @m.visit(
        m.ClassDef(
            bases=m.MatchIfTrue(
                lambda bases: any(
                    m.matches(
                        base.value,
                        m.Name("BaseModel")
                        | m.Attribute(attr=m.Name("BaseModel"))
                        | m.Name("Base")
                        | m.Attribute(attr=m.Name("Base")),
                    )
                    for base in bases
                )
            )
        )
    )
    def visit_basemodel(self, node: cst.ClassDef) -> None:
        if self.type_names and node.name.value not in self.type_names:
            return
        parents = self._bases_to_parents(node.bases)
        children = self._references_to_parents(typeVisitor._get_annassign_fields(node))
        basemodel_type = RawTypeMetadata(
            name=node.name.value,
            kind=ParameterTypeKind.CLASS,
            import_from=self.absolute_file_path,
            module_name=self._safe_name_to_module_name(node.name.value),
            project_path=self.project_path,
            is_leaf_DTO=True,
            is_core_module=self._is_current_file_core(),
            code=cst.Module([]).code_for_node(node),
            parents=frozenset(parents),
            attributes=frozenset(children),
        )
        self.type_mapping[basemodel_type.name] = basemodel_type
        self.visited_classes.add(id(node))

    @m.visit(m.ClassDef(bases=[m.Arg(value=m.Name("Enum") | m.Attribute(attr=m.Name("Enum")))]))
    def visit_enums(self, node: cst.ClassDef) -> None:
        if self.type_names and node.name.value not in self.type_names:
            return
        parents = self._bases_to_parents(node.bases)
        enum_type = RawTypeMetadata(
            name=node.name.value,
            kind=ParameterTypeKind.ENUM,
            import_from=self.absolute_file_path,
            module_name=self._safe_name_to_module_name(node.name.value),
            project_path=self.project_path,
            is_leaf_DTO=True,
            is_core_module=self._is_current_file_core(),
            code=cst.Module([]).code_for_node(node),
            parents=frozenset(parents),
            attributes=frozenset(),
        )
        self.type_mapping[enum_type.name] = enum_type
        self.visited_classes.add(id(node))

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        if self.type_names and node.name.value not in self.type_names:
            return False

        if id(node) in self.visited_classes:
            return False

        children = self._references_to_parents(typeVisitor._get_annassign_fields(node))
        parents = self._bases_to_parents(node.bases)
        class_type = RawTypeMetadata(
            name=node.name.value,
            kind=ParameterTypeKind.CLASS,
            import_from=self.absolute_file_path,
            module_name=self._safe_name_to_module_name(node.name.value),
            project_path=self.project_path,
            is_leaf_DTO=False,
            is_core_module=self._is_current_file_core(),
            code=cst.Module([]).code_for_node(node),
            parents=frozenset(parents),
            attributes=frozenset(children),
        )
        self.type_mapping[class_type.name] = class_type
        self.visited_classes.add(id(node))
        return None
