from __future__ import annotations

from typing import TYPE_CHECKING

from libcst import CSTTransformer, ImportAlias, ImportFrom, MetadataWrapper, Name, SimpleStatementLine

from core.languages.python.cst_utils import DepthTrackingMixin, build_module_path, parse_module_to_cst

if TYPE_CHECKING:
    from libcst import BaseStatement, ClassDef, FlattenSentinel, FunctionDef, Module, RemovalSentinel

    from aiservice.models.functions_to_optimize import FunctionToOptimize


def find_and_replace_function(module_cst: Module, function_name: str, new_node: BaseStatement) -> Module:
    """Find the first occurrence of a function definition with a given name in a libCST Module and replace it
    with a given libCST node.

    :param module_cst: A libCST Module object representing the parsed module.
    :param function_name: A string representing the name of the function to be replaced.
    :param new_node: A libCST node representing the new code to replace the function definition with.
    :return: A libCST Module object representing the modified module.
    """  # noqa: D205

    class FunctionReplacer(CSTTransformer):
        def __init__(self, name: str, node: BaseStatement) -> None:
            super().__init__()
            self.function_name = name
            self.new_node = node
            self.found = False

        def leave_FunctionDef(
            self, original_node: FunctionDef, updated_node: FunctionDef
        ) -> BaseStatement | FlattenSentinel[BaseStatement] | RemovalSentinel:
            if not self.found and original_node.name.value == self.function_name:
                self.found = True
                return self.new_node
            return updated_node

    wrapper = MetadataWrapper(module_cst, unsafe_skip_copy=True)
    replacer = FunctionReplacer(function_name, new_node)
    return wrapper.visit(replacer)


def replace_first_function_definition(module_str: str, function_name: str, stmt: str) -> str:
    """Replace the first occurrence of a function definition with a given name in a module string
    with a specified statement.

    :param module_str: A string representing the code of a Python module.
    :param function_name: A string representing the name of the function to be replaced.
    :param stmt: A string representing the statement to replace the function definition with.
    :return: A string representing the edited version of the input module string.
    """  # noqa: D205
    modified_module_cst = find_and_replace_function(
        parse_module_to_cst(module_str), function_name, parse_module_to_cst(stmt)
    )
    return modified_module_cst.code


class _ImportInserter(DepthTrackingMixin, CSTTransformer):
    """Transformer that replaces a function or class definition with an import statement."""

    def __init__(self, target_name: str, import_stmt: SimpleStatementLine) -> None:
        DepthTrackingMixin.__init__(self)
        CSTTransformer.__init__(self)
        self.target_name = target_name
        self.import_stmt = import_stmt
        self.replaced_count = 0

    def visit_ClassDef(self, node: ClassDef) -> None:  # noqa: ARG002
        self._visit_class()

    def leave_ClassDef(
        self, original_node: ClassDef, updated_node: ClassDef
    ) -> BaseStatement | FlattenSentinel[BaseStatement] | RemovalSentinel:
        if original_node.name.value == self.target_name and self._class_depth == 1 and self.replaced_count == 0:
            self.replaced_count += 1
            self._leave_class()
            return self.import_stmt
        self._leave_class()
        return updated_node

    def leave_FunctionDef(
        self, original_node: FunctionDef, updated_node: FunctionDef
    ) -> BaseStatement | FlattenSentinel[BaseStatement] | RemovalSentinel:
        if original_node.name.value == self.target_name and self._class_depth == 0 and self.replaced_count == 0:
            self.replaced_count += 1
            return self.import_stmt
        return updated_node


def replace_definition_with_import(module: Module, function: FunctionToOptimize, module_path: str) -> Module:
    """Replace a function or class definition with an import statement.

    Parameters
    ----------
    module : Module
        The CST module representing the test code.
    function : FunctionToOptimize
        The function to replace with an import.
    module_path : str
        Full module path of the function or class being replaced.

    Returns
    -------
    Module
        The modified CST module.

    """
    if not function.parents:
        new_import = SimpleStatementLine(
            body=[
                ImportFrom(
                    module=build_module_path(module_path), names=[ImportAlias(name=Name(value=function.function_name))]
                )
            ]
        )
        transformer = _ImportInserter(function.function_name, new_import)
        return module.visit(transformer)
    if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
        new_import = SimpleStatementLine(
            body=[
                ImportFrom(
                    module=build_module_path(module_path),
                    names=[ImportAlias(name=Name(value=function.parents[0].name))],
                )
            ]
        )
        transformer = _ImportInserter(function.parents[0].name, new_import)
        return module.visit(transformer)
    return module
