import logging
from uuid import UUID

from py_agent.models.parameter import KnownTypeMetadata, RawTypeMetadata
from py_agent.models.parameter_type_kind import ParameterTypeKind
from py_agent.models.type_metadata import TypeMetadata
from py_agent.services.dependency_extractor.dependency_info import DependencyInfo
from py_agent.services.dependency_extractor.type_metadata_converter import TypeMetadataConverter
from py_agent.services.type_metadata_extractor import TypeMetadataExtractor

logger = logging.getLogger(__name__)

DEPENDENCY_EXTRACTION_DEPTH = 4


class DependencyCollector:
    def __init__(self, extraction_depth: int = DEPENDENCY_EXTRACTION_DEPTH):
        self.extraction_depth = extraction_depth

    def collect_dependencies(
        self, type_metadata: TypeMetadata, raw_type: RawTypeMetadata, parameter_id: UUID
    ) -> list[DependencyInfo]:
        known_type, dep_list = self._get_deps_of_type(raw_type)
        result = [self._create_dep_info(x, parameter_id) for x in dep_list]
        result.append(self._create_dep_info(known_type, parameter_id, type_metadata))
        return result

    def _create_dep_info(
        self,
        known_type: KnownTypeMetadata,
        parameter_id: UUID,
        type_metadata: TypeMetadata | None = None,
    ) -> DependencyInfo:
        tm = (
            type_metadata if type_metadata else TypeMetadataConverter.convert_known_type(known_type)
        )
        return DependencyInfo(
            is_known=known_type.is_known,
            type_metadata=tm,
            code=known_type.code,
            is_data_object=known_type.is_DTO,
            parameter_id=parameter_id,
            nested_level=known_type.level,
            code_comments=None,
        )

    def _get_deps_of_type(
        self, type: RawTypeMetadata
    ) -> tuple[KnownTypeMetadata, list[KnownTypeMetadata]]:
        visited: dict[RawTypeMetadata, KnownTypeMetadata] = {}
        return self._resolve_type_depth(type, 0, visited)

    def _resolve_type_depth(
        self,
        type: RawTypeMetadata,
        current_depth: int,
        visited: dict[RawTypeMetadata, KnownTypeMetadata],
    ) -> tuple[KnownTypeMetadata, list[KnownTypeMetadata]]:
        if type in visited:
            return (
                KnownTypeMetadata(
                    **visited[type].model_dump(exclude={"level"}) | {"level": current_depth}
                ),
                [],
            )

        is_primitive_known = (
            type.kind == ParameterTypeKind.ENUM or type.kind == ParameterTypeKind.PRIMITIVE
        )
        if current_depth >= self.extraction_depth or is_primitive_known:
            is_dto_known = type.is_leaf_DTO and not type.attributes
            return KnownTypeMetadata(
                name=type.name,
                import_from=type.import_from,
                module_name=type.module_name,
                is_core_module=type.is_core_module,
                is_known=is_primitive_known or is_dto_known,
                is_DTO=type.is_leaf_DTO,
                kind=type.kind,
                code=type.code,
                level=current_depth,
            ), []

        resolved_parents_subtree: list[KnownTypeMetadata] = []
        resolved_parents: list[KnownTypeMetadata] = []
        for p in type.parents:
            new_parent = TypeMetadataExtractor.resolve_parent(p)
            if not new_parent:
                logger.debug(f"did not resolve {p}")
                continue
            known_obj, known_parents = self._resolve_type_depth(
                new_parent, current_depth + 1, visited
            )
            resolved_parents.append(known_obj)
            resolved_parents_subtree.append(known_obj)
            resolved_parents_subtree.extend(known_parents)

        resolved_types_subtree: list[KnownTypeMetadata] = []
        resolved_types: list[KnownTypeMetadata] = []
        for child in type.attributes:
            new_child = TypeMetadataExtractor.resolve_parent(child)
            if not new_child:
                logger.debug(f"did not resolve {child}")
                continue
            resolved_child, subtree_list = self._resolve_type_depth(
                new_child, current_depth + 1, visited
            )
            resolved_types.append(resolved_child)
            resolved_types_subtree.append(resolved_child)
            resolved_types_subtree.extend(subtree_list)

        resolved_types_subtree.extend(resolved_parents_subtree)
        all_attributes_known = all(x.is_known for x in [*resolved_types, *resolved_parents])
        is_dto = type.is_leaf_DTO or any(x.is_DTO for x in resolved_parents)

        resolved_type = KnownTypeMetadata(
            name=type.name,
            import_from=type.import_from,
            module_name=type.module_name,
            is_core_module=type.is_core_module,
            is_known=all_attributes_known and is_dto,
            is_DTO=is_dto,
            code=type.code,
            kind=type.kind,
            level=current_depth,
        )

        visited[type] = resolved_type
        return resolved_type, resolved_types_subtree
