from __future__ import annotations

from typing import cast

from aiservice.models.functions_to_optimize import FunctionToOptimize
from libcst import (
    Attribute,
    BaseStatement,
    ClassDef,
    CSTNode,
    CSTTransformer,
    FlattenSentinel,
    FunctionDef,
    ImportAlias,
    ImportFrom,
    MetadataWrapper,
    Module,
    Name,
    RemovalSentinel,
    parse_expression,
    parse_module,
)


def parse_module_to_cst(module_str: str) -> Module:
    """Parse a module string into its libCST representation.

    :param module_str: A string representing the code of a Python module.
    :return: A libCST Module object representing the parsed module.
    """
    return parse_module(module_str)


def find_and_replace_function(module_cst: Module, function_name: str, new_node: CSTNode) -> 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.
    """

    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) -> CSTNode:
            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 convert_cst_to_string(module_cst: Module) -> str:
    """Convert a libCST Module object into a string representation of the code.

    :param module_cst: A libCST Module object representing the parsed module.
    :return: A string representing the code of the module.
    """
    return module_cst.code


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.
    """
    modified_module_cst = find_and_replace_function(
        parse_module_to_cst(module_str), function_name, parse_module_to_cst(stmt)
    )
    return convert_cst_to_string(modified_module_cst)


def replace_definition_with_import(source_code: str, function: FunctionToOptimize, module_path: str) -> str:
    """Replace a function or class definition with an import statement in the given CST tree.

    Parameters
    ----------
    - tree (CSTModule): The CST tree representing the module.
    - function (FunctionToOptimize): The function to replace with an import.
    - module_path (str): Full module path of the function or class being replaced

    Returns
    -------
    - CSTModule: The modified CST tree.

    """

    class ImportInserter(CSTTransformer):
        def __init__(self, target_name: str, import_stmt: ImportFrom) -> None:
            super().__init__()
            self.target_name = target_name
            self.import_stmt = import_stmt

        def leave_ClassDef(
            self, original_node: ClassDef, updated_node: ClassDef
        ) -> BaseStatement | FlattenSentinel[BaseStatement] | RemovalSentinel:
            if original_node.name.value == self.target_name:
                return self.import_stmt
            return updated_node

        def leave_FunctionDef(
            self, original_node: FunctionDef, updated_node: FunctionDef
        ) -> FunctionDef | ImportFrom | FlattenSentinel[BaseStatement] | RemovalSentinel:
            if original_node.name.value == self.target_name:
                return self.import_stmt
            return updated_node

    tree = parse_module(source_code)
    if not function.parents:
        new_import = ImportFrom(
            module=cast(Attribute | Name | None, parse_expression(module_path)),
            names=[ImportAlias(name=Name(value=function.function_name))],
        )
        transformer = ImportInserter(function.function_name, new_import)
        return tree.visit(transformer).code
    if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
        new_import = ImportFrom(
            module=cast(Attribute | Name | None, parse_expression(module_path)),
            names=[ImportAlias(name=Name(value=function.parents[0].name))],
        )
        transformer = ImportInserter(function.parents[0].name, new_import)
        return tree.visit(transformer).code
    return tree.code
