import libcst as cst
from libcst import MetadataWrapper
from libcst.metadata import ByteSpanPositionProvider, PositionProvider

from py_agent.models.base import CustomBaseModel
from py_agent.models.class_metadata import ConstructorMetadata
from py_agent.models.parameter import Parameter
from py_agent.models.testable import Position


def get_current_class_name(current_class_node: cst.ClassDef | None) -> str | None:
    if current_class_node is None:
        return None
    return current_class_node.name.value


class ExtractedData(CustomBaseModel):
    function_code: str
    position: Position
    parameters: list[Parameter]
    return_parameter: Parameter | None = None
    class_code: str | None = None
    is_async: bool
    is_static: bool


class TestedCodeVisitor(cst.CSTVisitor):
    METADATA_DEPENDENCIES = (PositionProvider, ByteSpanPositionProvider)

    def __init__(
        self,
        module: cst.Module,
        metadata_wrapper: MetadataWrapper,
        absolute_file_path: str,
        relative_file_path: str,
        project_path: str,
        function_name: str,
        class_name: str | None,
    ):
        super().__init__()
        self.absolute_file_path = absolute_file_path
        self.relative_file_path = relative_file_path
        self.project_path = project_path
        self.function_name = function_name
        self.module = module
        self.metadata_wrapper = metadata_wrapper
        self.class_name = class_name
        self.extracted_data: ExtractedData | None = None
        self.is_class_extracted = False
        self.current_class_node: cst.ClassDef | None = None

    @property
    def is_extracted(self) -> bool:
        return self.extracted_data is not None

    @staticmethod
    def is_static_method(function_node: cst.FunctionDef) -> bool:
        for decorator in function_node.decorators:
            deco_expr = decorator.decorator
            if isinstance(deco_expr, cst.Name) and deco_expr.value == "staticmethod":
                return True
            if isinstance(deco_expr, cst.Attribute):
                if isinstance(deco_expr.attr, cst.Name) and deco_expr.attr.value == "staticmethod":
                    return True
        return False

    def visit_FunctionDef(self, function_node: cst.FunctionDef) -> bool | None:
        if self.is_extracted:
            return False

        if function_node.name.value != self.function_name:
            return True

        class_code = None
        class_name = get_current_class_name(self.current_class_node)

        if class_name == self.class_name:
            if self.class_name is not None:
                assert self.current_class_node is not None
                class_code = self.module.code_for_node(self.current_class_node)
                self.is_class_extracted = True
        else:
            return True

        function_code = self.module.code_for_node(function_node)
        position_metadata = self.get_metadata(PositionProvider, function_node)
        bytespan_metadata = self.get_metadata(ByteSpanPositionProvider, function_node)
        position = Position.from_metadata(position_metadata, bytespan_metadata)

        from py_agent.services.parameter_metadata_extractor import ParameterMetadataExtractor

        parameters = [
            ParameterMetadataExtractor(
                param.name.value, param.annotation, self.absolute_file_path, self.project_path
            ).extract()
            for param in function_node.params.params
        ]

        return_param = None
        if function_node.returns:
            return_param = ParameterMetadataExtractor(
                "", function_node.returns, self.absolute_file_path, self.project_path
            ).extract()

        is_async = bool(function_node.asynchronous)
        is_static = TestedCodeVisitor.is_static_method(function_node)

        self.extracted_data = ExtractedData(
            function_code=function_code,
            position=position,
            parameters=parameters,
            return_parameter=return_param,
            class_code=class_code,
            is_async=is_async,
            is_static=is_static,
        )
        return False

    def visit_ClassDef(self, class_node: cst.ClassDef) -> bool | None:
        self.current_class_node = class_node
        return not self.is_extracted

    def leave_ClassDef(self, class_node: cst.ClassDef) -> None:
        self.current_class_node = None

    def get_constructor_metadata(self) -> ConstructorMetadata | None:
        if self.class_name is None or not self.is_extracted:
            return None

        visitor = TestedCodeVisitor(
            module=self.module,
            metadata_wrapper=self.metadata_wrapper,
            absolute_file_path=self.absolute_file_path,
            relative_file_path=self.relative_file_path,
            project_path=self.project_path,
            function_name="__init__",
            class_name=self.class_name,
        )
        self.metadata_wrapper.visit(visitor)
        extracted_data = visitor.extracted_data
        if not extracted_data:
            return None
        return ConstructorMetadata(
            code=extracted_data.function_code, parameters=extracted_data.parameters
        )
