import logging
import sys
from pathlib import Path

import libcst as cst

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.services.import_jedi_extractor import ImportJediExtractor, JediMetadata
from py_agent.utils.constants import PRIMITIVE_TYPES

logger = logging.getLogger(__name__)

# Thread-local resolved types cache — cleared per extraction call
_resolved_types: dict[RawParentTypeMetadata, RawTypeMetadata] = {}


def get_resolved_types_cache() -> dict[RawParentTypeMetadata, RawTypeMetadata]:
    return _resolved_types


def clear_resolved_types_cache() -> None:
    global _resolved_types
    _resolved_types = {}


def _read_file(file_path: Path) -> str:
    with open(file_path, encoding="utf-8") as f:
        return f.read()


class TypeMetadataExtractor:
    def __init__(self, absolute_file_path: str, project_path: str):
        self.absolute_file_path = absolute_file_path
        self.project_path = project_path

    @staticmethod
    def is_stdlib(path: Path) -> bool:
        return path.is_relative_to(Path(sys.base_prefix))

    @staticmethod
    def resolve_parent(parent: RawParentTypeMetadata) -> RawTypeMetadata | None:
        if parent.is_core_module or TypeMetadataExtractor.is_stdlib(parent.import_from):
            return None

        cache = get_resolved_types_cache()
        if parent in cache:
            return cache[parent]

        from py_agent.visitors.type_visitor import typeVisitor

        code = _read_file(parent.import_from)
        module = cst.parse_module(code)
        wrapper = cst.MetadataWrapper(module)
        visitor = typeVisitor(
            absolute_file_path=str(parent.import_from),
            project_path=str(parent.project_path),
            type_names=[parent.name],
        )
        wrapper.visit(visitor)
        type_mapping = visitor.get_type_mapping()
        if parent.name not in type_mapping:
            logger.error(f"jedi could not resolve parent: {parent}")
            return None
        cache[parent] = type_mapping[parent.name]
        return type_mapping[parent.name]

    @staticmethod
    def _recursive_unravel(node: cst.CSTNode) -> set[AnnotationType]:
        if isinstance(node, cst.Name):
            name = node.value
            kind = ParameterTypeKind.PRIMITIVE if name in PRIMITIVE_TYPES else ParameterTypeKind.ANY
            return {AnnotationType(name, kind, namespace=None, full_name=name)}

        if isinstance(node, cst.Attribute):
            full_name = cst.Module([]).code_for_node(node)
            namespace = cst.Module([]).code_for_node(node.value)
            name = node.attr.value
            kind = ParameterTypeKind.PRIMITIVE if name in PRIMITIVE_TYPES else ParameterTypeKind.ANY
            return {AnnotationType(name, kind, namespace=namespace, full_name=full_name)}

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

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

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

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

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

        return set()

    @staticmethod
    def recursive_unravel(node: cst.CSTNode) -> set[AnnotationType]:
        return TypeMetadataExtractor._recursive_unravel(node)

    def extract(self, annotation: cst.CSTNode | None) -> list[RawTypeMetadata]:
        if annotation is None:
            return []

        jedi_extractor = ImportJediExtractor(
            absolute_file_path=self.absolute_file_path,
            project_path=self.project_path,
        )

        annotation_type_list = TypeMetadataExtractor.recursive_unravel(annotation)

        groupings: dict[Path, list[JediMetadata]] = {}
        metadata = [
            (annotation_type, jedi_extractor.get_jedi_metadata(annotation_type.full_name))
            for annotation_type in annotation_type_list
            if annotation_type.kind != ParameterTypeKind.PRIMITIVE
        ]

        result: list[RawTypeMetadata] = []
        for annotation_type, jedi_metadata in metadata:
            if not jedi_metadata:
                logger.debug(f"jedi could not find {annotation_type}")
                continue
            if not jedi_metadata.file_path or jedi_metadata.is_core:
                continue
            if jedi_metadata.file_path not in groupings:
                groupings[jedi_metadata.file_path] = []
            groupings[jedi_metadata.file_path].append(jedi_metadata)

        from py_agent.visitors.type_visitor import typeVisitor

        for path, metadata_list in groupings.items():
            code = _read_file(path)
            module = cst.parse_module(code)
            wrapper = cst.MetadataWrapper(module)
            visitor = typeVisitor(
                absolute_file_path=str(path),
                project_path=self.project_path,
                type_names=[x.name for x in metadata_list],
            )
            wrapper.visit(visitor)
            type_mapping = visitor.get_type_mapping()
            for jm in metadata_list:
                if jm.name not in type_mapping:
                    logger.debug(f"error with type visiting and {jm.name}")
                    continue
                result.append(type_mapping[jm.name])

        return result
