from pathlib import Path
from typing import Any

from py_agent.models.code_dependency import CodeDependency, ParentParameterRelations
from py_agent.models.enum_metadata import EnumMetadata
from py_agent.models.missing_dependency import DependencyMetadata, MissingDependency
from py_agent.models.parameter import Parameter
from py_agent.models.parameter_type_kind import ParameterTypeKind
from py_agent.models.type_metadata import TypeMetadata
from py_agent.services.dependency_extractor.dependency_collector import DEPENDENCY_EXTRACTION_DEPTH
from py_agent.services.dependency_extractor.dependency_info import DependencyInfo
from py_agent.services.dependency_extractor.dependency_info_extractor import DependencyInfoExtractor
from py_agent.utils.list_utils import flatten, uniq_with


def get_contained_in_tested_file(relative_file_path: str, type_metadata: TypeMetadata) -> bool:
    if not type_metadata.is_core_module:
        return False
    return type_metadata.import_path in relative_file_path


class DependencyExtractor:
    def __init__(
        self, relative_file_path: str, extraction_depth: int = DEPENDENCY_EXTRACTION_DEPTH
    ):
        self.relative_file_path = relative_file_path
        self.extraction_depth = extraction_depth

    def get_raw_dependencies_results(
        self, parameter: Parameter
    ) -> tuple[list[CodeDependency], list[MissingDependency]]:
        """Resolve dependencies for a single parameter's type tree.

        Walks the parameter's type annotations up to ``extraction_depth``
        levels using :class:`DependencyInfoExtractor`, then partitions the
        results into *known* (source available, non-primitive, non-stdlib)
        and *missing* (top-level types whose source could not be resolved).

        Args:
            parameter: A function parameter (or return type) whose type
                annotations will be traversed for dependency discovery.

        Returns:
            A ``(known, missing)`` tuple of dependency lists before
            deduplication — callers are expected to reduce/unique them.
        """
        dependencies_info = DependencyInfoExtractor.get_dependencies_info(
            parameter, self.extraction_depth
        )
        return self.get_known_dependencies(dependencies_info), self.get_missing_dependencies(
            dependencies_info
        )

    def get_known_dependencies(
        self, dependencies_info: list[DependencyInfo]
    ) -> list[CodeDependency]:
        known = [
            dep
            for dep in dependencies_info
            if dep.is_known
            and not dep.type_metadata.is_core_module
            and not dep.type_metadata.is_primitive
            and dep.code
        ]
        return [self.dependency_info_to_code_dependency(d) for d in known]

    def get_missing_dependencies(
        self, dependencies_info: list[DependencyInfo]
    ) -> list[MissingDependency]:
        missing = [dep for dep in dependencies_info if not dep.is_known and dep.nested_level == 0]
        return [self.dependency_info_to_missing_dependency(d) for d in missing]

    def dependency_predicate(
        self, d1: CodeDependency | MissingDependency, d2: CodeDependency | MissingDependency
    ) -> bool:
        return d1.name == d2.name and d1.kind == d2.kind

    def reduce_known_dependencies(self, known_deps: list[CodeDependency]) -> list[CodeDependency]:
        collected: list[CodeDependency] = []
        for dependency in known_deps:
            found = next((d for d in collected if self.dependency_predicate(d, dependency)), None)
            if found is None:
                collected.append(dependency)
                continue
            if dependency.parent_parameter_ids:
                found.parent_parameter_ids.extend(dependency.parent_parameter_ids)
            if dependency.parent_parameter_relations:
                found.parent_parameter_relations.extend(dependency.parent_parameter_relations)
        return collected

    def get_unique_dependencies(
        self, dependencies: list[MissingDependency]
    ) -> list[MissingDependency]:
        return uniq_with(dependencies, self.dependency_predicate)

    def dependency_info_to_code_dependency(self, dependency_info: DependencyInfo) -> CodeDependency:
        parent_parameter_ids = (
            [dependency_info.parameter_id] if dependency_info.parameter_id else []
        )
        parent_parameter_relations = (
            [
                ParentParameterRelations(
                    parameterParentId=dependency_info.parameter_id,
                    depthLevelInParent=dependency_info.nested_level,
                )
            ]
            if dependency_info.parameter_id
            else []
        )
        return CodeDependency(
            name=dependency_info.type_metadata.name,
            kind=dependency_info.type_metadata.kind,
            code=dependency_info.code,
            importPath=dependency_info.type_metadata.import_path,
            moduleName=dependency_info.type_metadata.module_name,
            containedInTestedFile=get_contained_in_tested_file(
                self.relative_file_path, dependency_info.type_metadata
            ),
            isDataObject=dependency_info.is_data_object,
            parentParameterIds=parent_parameter_ids,
            parentParameterRelations=parent_parameter_relations,
        )

    def dependency_info_to_missing_dependency(
        self, dependency_info: DependencyInfo
    ) -> MissingDependency:
        parent_parameter_ids = (
            [dependency_info.parameter_id] if dependency_info.parameter_id else []
        )
        return MissingDependency(
            name=dependency_info.type_metadata.name,
            kind=dependency_info.type_metadata.kind,
            code=dependency_info.code,
            importPath=dependency_info.type_metadata.import_path,
            containedInTestedFile=get_contained_in_tested_file(
                self.relative_file_path, dependency_info.type_metadata
            ),
            isDataObject=dependency_info.is_data_object,
            parentParameterIds=parent_parameter_ids,
            isCoreModule=dependency_info.type_metadata.is_core_module,
            codeComments=dependency_info.code_comments,
        )

    def extract(
        self,
        params: list[Parameter],
        external_functions: list[tuple[Any, ...]],
        enums: list[EnumMetadata],
    ) -> tuple[list[CodeDependency], list[MissingDependency]]:
        """Extract all dependencies for a function under test.

        Walks each parameter's type tree (up to ``extraction_depth`` levels) to
        discover the types it depends on, then separates them into *known*
        dependencies (source code is available) and *missing* ones (source is
        unavailable or lives in an external package).  External function calls
        and enum types referenced by the function are folded in as well.

        Args:
            params: Parameters of the function under test (including return type).
            external_functions: Pairs of ``(ExternalMethodInfo, extracted_data)``
                for every non-method function called by the tested function.
            enums: Enum types referenced in the function's body or signature.

        Returns:
            A ``(known, missing)`` tuple where *known* contains dependencies
            whose source code was resolved and *missing* lists those that
            could not be resolved (e.g. third-party or stdlib types).
        """
        individual_results = [self.get_raw_dependencies_results(p) for p in params]

        known_deps_duplicates = flatten([known for known, _ in individual_results])
        known_deps_duplicates.extend(self.extract_enums(enums))
        known_deps = self.reduce_known_dependencies(known_deps_duplicates)

        missing_deps_duplicates = flatten([missing for _, missing in individual_results])
        missing_deps_duplicates.extend(self.extract_external_methods(external_functions))
        missing_deps = self.get_unique_dependencies(missing_deps_duplicates)

        return known_deps, missing_deps

    def extract_external_methods(
        self, external_methods: list[tuple[Any, ...]]
    ) -> list[MissingDependency]:
        result = []
        for dep, extracted_data in external_methods:
            missing_dep = MissingDependency(
                name=dep.name,
                code=extracted_data.code,
                containedInTestedFile=False,
                importPath=Path(dep.jedi_metadata.file_path).as_posix()
                if dep.jedi_metadata.file_path
                else None,
                isCoreModule=dep.jedi_metadata.is_core,
                kind=ParameterTypeKind.FUNCTION,
                isDataObject=False,
                parentParameterIds=[],
                codeComments=None,
                dependencyMetadata=DependencyMetadata(
                    classMetadata=extracted_data.method_class,
                    methodsMetadata=[extracted_data],
                ),
            )
            result.append(missing_dep)
        return result

    def extract_enums(self, enums: list[EnumMetadata]) -> list[CodeDependency]:
        return [
            CodeDependency(
                name=enum.name,
                importPath=Path(enum.module_path).resolve().as_posix(),
                moduleName=enum.module_name,
                containedInTestedFile=False,
                kind=ParameterTypeKind.ENUM,
                code=enum.code,
                isDataObject=True,
                parentParameterIds=[],
                parentParameterRelations=[],
            )
            for enum in enums
        ]
