from __future__ import annotations

from typing import TypeVar, cast

from aiservice.models.functions_to_optimize import FunctionToOptimize
from libcst import (
    Assert,
    Assign,
    AssignTarget,
    Attribute,
    BaseSmallStatement,
    Call,
    Comment,
    Comparison,
    CSTTransformer,
    EmptyLine,
    Expr,
    FunctionDef,
    ImportAlias,
    ImportFrom,
    ImportStar,
    ListComp,
    Module,
    Name,
    Newline,
    Not,
    Pass,
    RemovalSentinel,
    RemoveFromParent,
    SimpleStatementLine,
    SimpleStatementSuite,
    UnaryOperation,
    matchers,
    parse_expression,
)

SimpleStatementType = TypeVar("SimpleStatementType", SimpleStatementLine, SimpleStatementSuite)


class RemoveAssertTransformer(CSTTransformer):
    def __init__(self, function: FunctionToOptimize, dependent_function_names: list[str]) -> None:
        self.function_object: FunctionToOptimize = function
        self.class_name: str | None = None
        self.only_function_name: str = function.function_name
        if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
            self.class_name = function.top_level_parent_name
        self.dependent_function_names: list[str] = dependent_function_names
        self.assignments: dict[str, Call] = {}

        # List of unittest assert methods to be removed
        self.unittest_assert_methods = [
            "assertEqual",
            "assertNotEqual",
            "assertTrue",
            "assertFalse",
            "assertIs",
            "assertIsNot",
            "assertIsNone",
            "assertIsNotNone",
            "assertIn",
            "assertNotIn",
            "assertIsInstance",
            "assertNotIsInstance",
            "assertRaises",
            "assertRaisesRegex",
            "assertWarns",
            "assertWarnsRegex",
            "assertLogs",
            "assertNoLogs",
            "assertAlmostEqual",
            "assertNotAlmostEqual",
            "assertGreater",
            "assertGreaterEqual",
            "assertLess",
            "assertLessEqual",
            "assertRegex",
            "assertNotRegex",
            "assertCountEqual",
            "assertMultiLineEqual",
            "assertSequenceEqual",
            "assertListEqual",
            "assertTupleEqual",
            "assertSetEqual",
            "assertDictEqual",
        ]

    def leave_ImportFrom(self, original_node: ImportFrom, updated_node: ImportFrom) -> RemovalSentinel | ImportFrom:
        # Skip checks for wildcard imports
        if isinstance(updated_node.names, ImportStar):
            return updated_node
        # Check if the import should be removed because it matches function or class names
        for alias in updated_node.names:
            if alias.name.value in [self.only_function_name, self.class_name]:
                return RemoveFromParent()

        return updated_node

    def leave_Assign(self, original_node: Assign, updated_node: Assign) -> Assign:
        # Track the assignment if it is a function call
        if (
            isinstance(updated_node.value, Call)
            and isinstance(updated_node.targets[0].target, Name)
            and self.is_target_function_node(updated_node.value)
        ):
            var_name: str = updated_node.targets[0].target.value
            self.assignments[var_name] = updated_node.value
            return updated_node.with_changes(targets=[AssignTarget(target=Name("codeflash_output"))])
        return updated_node

    def leave_SimpleStatementLine(
        self, original_node: SimpleStatementLine, updated_node: SimpleStatementLine
    ) -> SimpleStatementLine | RemovalSentinel:
        if updated_node.body:
            # Handle multiple statements on one line
            if len(updated_node.body) > 1:
                new_body: list[BaseSmallStatement] = []
                for stmt in updated_node.body:
                    if isinstance(stmt, Assert):
                        new_stmt = self._handle_assert(stmt)
                        if isinstance(new_stmt, Assign):
                            new_body.append(new_stmt)
                    else:
                        new_body.append(stmt)

                # Clean up trailing semicolon
                if new_body and isinstance(new_body[-1], BaseSmallStatement) and hasattr(new_body[-1], "semicolon"):
                    new_body[-1] = new_body[-1].with_changes(semicolon=None)
                return updated_node.with_changes(body=new_body)

            # Handle single pytest assert statement
            if isinstance(updated_node.body[0], Assert):
                new_node = self._handle_assert(updated_node.body[0])
                return updated_node.with_changes(body=[new_node]) if new_node else RemoveFromParent()

            # Handle unittest assert
            if matchers.matches(
                updated_node,
                matchers.SimpleStatementLine(
                    body=[
                        matchers.Expr(
                            matchers.Call(func=matchers.Attribute(value=matchers.Name("self"), attr=matchers.Name()))
                        )
                    ]
                ),
            ):
                unittest_node = self._handle_unittest_assert(updated_node)
                return unittest_node if unittest_node else RemoveFromParent()

        return updated_node

    def leave_SimpleStatementSuite(
        self, original_node: SimpleStatementSuite, updated_node: SimpleStatementSuite
    ) -> SimpleStatementSuite:
        if updated_node.body:
            # Handle multiple statements on one line
            if len(updated_node.body) > 1:
                new_body: list[BaseSmallStatement] = []
                for stmt in updated_node.body:
                    if isinstance(stmt, Assert):
                        new_stmt = self._handle_assert(stmt)
                        if isinstance(new_stmt, Assign):
                            new_body.append(new_stmt)
                    else:
                        new_body.append(stmt)

                # Clean up trailing semicolon
                if new_body and isinstance(new_body[-1], BaseSmallStatement) and hasattr(new_body[-1], "semicolon"):
                    new_body[-1] = new_body[-1].with_changes(semicolon=None)
                return updated_node.with_changes(body=new_body if new_body else [Pass()])

            # Handle single pytest assert statement
            if isinstance(updated_node.body[0], Assert):
                new_node = self._handle_assert(updated_node.body[0])
                return updated_node.with_changes(body=[new_node if new_node else Pass()])

            # Handle unittest assert
            if matchers.matches(
                updated_node,
                matchers.SimpleStatementSuite(
                    body=[
                        matchers.Expr(
                            matchers.Call(func=matchers.Attribute(value=matchers.Name("self"), attr=matchers.Name()))
                        )
                    ]
                ),
            ):
                unittest_node = self._handle_unittest_assert(updated_node)
                return (
                    updated_node.with_changes(body=unittest_node)
                    if unittest_node
                    else updated_node.with_changes(body=[Pass()])
                )

        return updated_node

    def _handle_assert(self, assert_node: Assert) -> Assign | None:
        # Handle "assert function_call()" and "assert not function_call()" cases
        if (isinstance(assert_node.test, Call) and self.is_target_function_node(assert_node.test)) or (
            isinstance(assert_node.test, UnaryOperation)
            and isinstance(assert_node.test.operator, Not)
            and isinstance(assert_node.test.expression, Call)
            and self.is_target_function_node(assert_node.test.expression)
        ):
            return Assign(targets=[AssignTarget(target=Name("codeflash_output"))], value=assert_node.test)

        # Handle "assert function_call() == value" cases
        if isinstance(assert_node.test, Comparison):
            comparison: Comparison = assert_node.test
            for side in [comparison.left, comparison.comparisons[0].comparator]:
                if (
                    isinstance(side, ListComp)  # Check if it's a list comprehension
                    and isinstance(side.elt, Call)
                    and self.is_target_function_node(side.elt)
                ) or (isinstance(side, Call) and self.is_target_function_node(side)):
                    return Assign(targets=[AssignTarget(target=Name("codeflash_output"))], value=side)
        return None

    def _handle_unittest_assert(self, updated_node: SimpleStatementType) -> SimpleStatementType | None:
        if isinstance(updated_node.body[0], (Assign, Expr)) and isinstance(updated_node.body[0].value, Call):
            method_call = updated_node.body[0].value
            if not (
                isinstance(method_call.func, Attribute) and method_call.func.attr.value in self.unittest_assert_methods
            ):
                return updated_node

            # Check arguments for function calls
            for arg in method_call.args:
                if isinstance(arg.value, ListComp):
                    check_node = arg.value.elt
                else:
                    check_node = arg.value
                if isinstance(check_node, Call) and self.is_target_function_node(check_node):
                    new_assign = Assign(targets=[AssignTarget(target=Name("codeflash_output"))], value=arg.value)
                    return updated_node.with_changes(body=[new_assign])

        return None

    def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> RemovalSentinel | FunctionDef:
        if not updated_node.body.body:
            # Remove empty functions
            return RemoveFromParent()

        return updated_node

    def is_target_function_node(self, node: Call) -> bool:
        return (isinstance(node.func, Name) and node.func.value == self.only_function_name) or (
            isinstance(node.func, Attribute) and node.func.attr.value == self.only_function_name
        )

    def leave_Module(self, original_node: Module, updated_node: Module) -> Module:
        # Create a newline and comment
        footer = EmptyLine(
            indent=True,
            newline=Newline(),
            comment=Comment(
                "# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."
            ),
        )

        # Add footer to the end of the module
        return updated_node.with_changes(footer=[footer])


def remove_asserts_from_test(
    module: Module, function_to_optimize: FunctionToOptimize, helper_function_names: list[str], module_path: str
) -> Module:
    transformer = RemoveAssertTransformer(function_to_optimize, helper_function_names)
    modified_tree = module.visit(transformer)
    add_import = SimpleStatementLine(
        body=[
            ImportFrom(
                module=cast(Attribute | Name | None, parse_expression(module_path)),
                names=[ImportAlias(name=Name(function_to_optimize.top_level_parent_name))],
            )
        ]
    )
    new_body = [add_import] + list(modified_tree.body)
    modified_tree = modified_tree.with_changes(body=tuple(new_body))
    return modified_tree
