from __future__ import annotations

from typing import TYPE_CHECKING, Any, cast

from libcst import (
    Assert,
    Assign,
    AssignTarget,
    Attribute,
    BaseSmallStatement,
    Call,
    Comment,
    Comparison,
    CSTTransformer,
    EmptyLine,
    Expr,
    FlattenSentinel,
    ImportAlias,
    ImportFrom,
    ImportStar,
    ListComp,
    Name,
    Newline,
    Not,
    Pass,
    RemoveFromParent,
    SimpleStatementLine,
    SimpleStatementSuite,
    UnaryOperation,
    parse_expression,
    parse_module,
)

if TYPE_CHECKING:
    from libcst import FunctionDef, Module, RemovalSentinel

    from aiservice.models.functions_to_optimize import FunctionToOptimize


class StatementHandler:
    def __init__(self, transformer: RemoveAssertTransformer) -> None:
        self.transformer = transformer

    def handle_statement(
        self, updated_node: SimpleStatementLine | SimpleStatementSuite
    ) -> SimpleStatementLine | SimpleStatementSuite | RemovalSentinel | FlattenSentinel[BaseSmallStatement] | None:
        if not updated_node.body:
            # Empty body is malformed - handle it properly
            if isinstance(updated_node, SimpleStatementSuite):
                return updated_node.with_changes(body=[Pass()])
            # SimpleStatementLine with empty body should be removed
            return RemoveFromParent()

        if len(updated_node.body) == 1:
            return self._handle_single_statement(updated_node)
        return self._handle_multiple_statements(updated_node)

    def _handle_single_statement(
        self, updated_node: SimpleStatementLine | SimpleStatementSuite
    ) -> SimpleStatementLine | SimpleStatementSuite | RemovalSentinel | FlattenSentinel[BaseSmallStatement] | None:
        statement = updated_node.body[0]

        if isinstance(statement, Assert):
            return self._handle_assert_or_unittest(updated_node, statement)
        if isinstance(statement, Assign):
            return self._handle_assign(updated_node, statement)
        if self.transformer.is_unittest_external_assert_statement(statement):
            return self._handle_assert_or_unittest(updated_node, cast("Expr", statement))

        return updated_node

    def _handle_multiple_statements(
        self, updated_node: SimpleStatementLine | SimpleStatementSuite
    ) -> SimpleStatementLine | SimpleStatementSuite | RemovalSentinel | FlattenSentinel[BaseSmallStatement]:
        new_body: list[BaseSmallStatement] = []
        processed_codeflash_assignments = False

        for stmt in updated_node.body:
            assign_node = None
            if isinstance(stmt, Assert):
                assign_node = self.transformer.transform_assert_to_assignment(stmt)
            elif isinstance(stmt, Assign):
                processed_codeflash_assignments = self._process_assign_in_multiple(
                    stmt, new_body, processed_codeflash_assignments
                )
            elif self.transformer.is_unittest_external_assert_statement(stmt):
                assign_node = self.transformer.transform_unittest_to_assignment(cast("Expr", stmt))
            else:
                new_body.append(stmt)

            if assign_node:
                new_body.append(assign_node)
                processed_codeflash_assignments = True

        if new_body and hasattr(new_body[-1], "semicolon"):
            new_body[-1] = new_body[-1].with_changes(semicolon=None)

        if not new_body:
            if isinstance(updated_node, SimpleStatementSuite):
                new_body = [Pass()]
            else:
                # SimpleStatementLine with empty body is malformed - remove it entirely
                return RemoveFromParent()

        return updated_node.with_changes(body=new_body)

    def _process_assign_in_multiple(
        self,
        stmt: Assign,
        new_body: list[BaseSmallStatement],
        processed_flag: bool,  # noqa: FBT001
    ) -> bool:
        if self.transformer.is_redundant_codeflash_assign(stmt, processed_flag):
            return processed_flag

        assign_result = self.transformer.handle_assignment(stmt)
        processed = processed_flag

        if isinstance(assign_result, FlattenSentinel | Assign):
            nodes = assign_result if isinstance(assign_result, FlattenSentinel) else [assign_result]
            for node in nodes:
                if isinstance(node, Assign) and self.transformer.is_codeflash_output_assign(node):
                    if not processed:
                        new_body.append(node)
                        processed = True
                elif isinstance(node, BaseSmallStatement):
                    new_body.append(node)

        return processed

    def _handle_assert_or_unittest(
        self, updated_node: SimpleStatementLine | SimpleStatementSuite, node: Assert | Expr
    ) -> SimpleStatementLine | SimpleStatementSuite | RemovalSentinel | FlattenSentinel[BaseSmallStatement] | None:
        if isinstance(node, Assert):
            assign_node = self.transformer.transform_assert_to_assignment(node)
        else:
            assign_node = self.transformer.transform_unittest_to_assignment(node)

        if assign_node is None:
            return (
                RemoveFromParent()
                if isinstance(updated_node, SimpleStatementLine)
                else updated_node.with_changes(body=[Pass()])
            )

        if len(updated_node.body) > 1:
            if isinstance(updated_node, SimpleStatementLine):
                non_assert_stmts: list[BaseSmallStatement] = [
                    stmt
                    for stmt in updated_node.body
                    if not self.transformer.is_unittest_external_assert_statement(stmt)
                ]
                return updated_node.with_changes(body=[assign_node, *non_assert_stmts])

            result_lines = [SimpleStatementLine(body=[assign_node])]
            non_assert_stmts = [
                stmt for stmt in updated_node.body if not self.transformer.is_unittest_external_assert_statement(stmt)
            ]
            if non_assert_stmts:
                result_lines.extend([SimpleStatementLine(body=[stmt]) for stmt in non_assert_stmts])
            return FlattenSentinel(cast("list[BaseSmallStatement]", result_lines))

        return updated_node.with_changes(body=[assign_node])

    def _handle_assign(
        self, updated_node: SimpleStatementLine | SimpleStatementSuite, assign_node: Assign
    ) -> SimpleStatementLine | SimpleStatementSuite | FlattenSentinel[BaseSmallStatement]:
        assign_result = self.transformer.handle_assignment(assign_node)
        if isinstance(assign_result, Assign):
            return updated_node.with_changes(body=[assign_result])
        return assign_result


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] = {}
        self.codeflash_references: dict[str, str] = {}
        self.variable_reassignments: dict[str, set[str]] = {}

        self.statement_handler = StatementHandler(self)
        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",
        ]
        self.numpy_assert_methods = [
            "assert_",
            "assert_allclose",
            "assert_almost_equal",
            "assert_approx_equal",
            "assert_array_almost_equal",
            "assert_array_almost_equal_nulp",
            "assert_array_compare",
            "assert_array_equal",
            "assert_array_less",
            "assert_array_max_ulp",
            "assert_equal",
            "assert_no_gc_cycles",
            "assert_no_warnings",
            "assert_raises",
            "assert_raises_regex",
            "assert_string_equal",
            "assert_warns",
        ]
        self.pandas_assert_methods = [
            "assert_frame_equal",
            "assert_series_equal",
            "assert_index_equal",
            "assert_extension_array_equal",
        ]
        self.all_assert_methods = self.unittest_assert_methods + self.numpy_assert_methods + self.pandas_assert_methods
        self.valid_assert_prefixes = {"self", "numpy", "np", "pandas", "pd"}

    def leave_ImportFrom(self, original_node: ImportFrom, updated_node: ImportFrom) -> RemovalSentinel | ImportFrom:  # noqa: ARG002
        if isinstance(updated_node.names, ImportStar):
            return updated_node
        if all(alias.name.value in [self.only_function_name, self.class_name] for alias in updated_node.names):
            return RemoveFromParent()
        return updated_node

    def is_redundant_codeflash_assign(self, stmt: Assign, processed_flag: bool) -> bool:  # noqa: FBT001
        return (
            processed_flag
            and isinstance(stmt.targets[0].target, Name)
            and stmt.targets[0].target.value == "codeflash_output"
        )

    def is_codeflash_output_assign(self, node: Assign) -> bool:
        return isinstance(node.targets[0].target, Name) and node.targets[0].target.value == "codeflash_output"

    def is_unittest_external_assert_statement(self, stmt: BaseSmallStatement) -> bool:
        if not (
            isinstance(stmt, Expr)
            and isinstance(stmt.value, Call)
            and isinstance(stmt.value.func, Attribute)
            and stmt.value.func.attr.value in self.all_assert_methods
        ):
            return False
        func_value = stmt.value.func.value
        if isinstance(func_value, Name):
            return func_value.value in self.valid_assert_prefixes
        if not isinstance(func_value, Attribute):
            return False
        if isinstance(func_value.value, Name):
            return func_value.value.value in self.valid_assert_prefixes
        if isinstance(func_value.value, Attribute) and isinstance(func_value.value.value, Name):
            return func_value.value.value.value in self.valid_assert_prefixes
        return False

    def transform_assert_to_assignment(self, assert_node: Assert) -> Assign | None:
        if self.is_direct_function_call_assert(assert_node):
            return self.create_codeflash_assignment(assert_node.test)

        if isinstance(assert_node.test, Comparison):
            return self.handle_comparison_assert(assert_node.test)

        return None

    def transform_unittest_to_assignment(self, unittest_assert_node: Expr) -> Assign | None:
        method_call = cast("Call", unittest_assert_node.value)
        for arg in method_call.args:
            check_node = cast("BaseSmallStatement", arg.value)
            if self._should_create_assignment_from_check_node(check_node):
                return self.create_codeflash_assignment(check_node)
        return None

    def _should_create_assignment_from_check_node(self, check_node: BaseSmallStatement) -> bool:
        if isinstance(check_node, ListComp):
            return isinstance(check_node.elt, Call) and self.is_target_function_node(check_node.elt)
        return (
            self.contains_codeflash_reference(check_node)
            and isinstance(check_node, Call)
            and self.is_target_function_node(check_node)
        )

    def handle_assignment(self, updated_node: Assign) -> Assign | FlattenSentinel[BaseSmallStatement]:
        if not isinstance(updated_node.targets[0].target, Name):
            return updated_node

        var_name: str = updated_node.targets[0].target.value

        if self._is_self_codeflash_assignment(updated_node, var_name):
            return updated_node

        if isinstance(updated_node.value, Call) and self.is_target_function_node(updated_node.value):
            return self._handle_target_function_assignment(updated_node, var_name)

        if self._is_codeflash_output_assignment(updated_node, var_name) and (
            var_name not in self.codeflash_references or self.codeflash_references[var_name] != "codeflash_output"
        ):
            self.codeflash_references[var_name] = "codeflash_output"

        if self._is_reference_chain_assignment(updated_node):
            return self._handle_reference_chain_assignment(updated_node, var_name)

        return updated_node

    def _is_self_codeflash_assignment(self, updated_node: Assign, var_name: str) -> bool:
        return (
            var_name == "codeflash_output"
            and isinstance(updated_node.value, Name)
            and updated_node.value.value == "codeflash_output"
        )

    def _is_codeflash_output_assignment(self, updated_node: Assign, _var_name: str) -> bool:
        return isinstance(updated_node.value, Name) and updated_node.value.value == "codeflash_output"

    def _is_reference_chain_assignment(self, updated_node: Assign) -> bool:
        return isinstance(updated_node.value, Name) and updated_node.value.value in self.codeflash_references

    def _handle_target_function_assignment(
        self, updated_node: Assign, var_name: str
    ) -> FlattenSentinel[BaseSmallStatement]:
        self.assignments[var_name] = cast("Call", updated_node.value)
        self.codeflash_references[var_name] = "codeflash_output"
        codeflash_assign = updated_node.with_changes(targets=[AssignTarget(target=Name("codeflash_output"))])
        original_var_assign = Assign(targets=[AssignTarget(target=Name(var_name))], value=Name("codeflash_output"))
        return FlattenSentinel([codeflash_assign, original_var_assign])

    def _handle_reference_chain_assignment(self, updated_node: Assign, var_name: str) -> Assign:
        source_var = cast("Name", updated_node.value).value
        self.codeflash_references[var_name] = source_var
        if source_var not in self.variable_reassignments:
            self.variable_reassignments[source_var] = set()
        self.variable_reassignments[source_var].add(var_name)
        return updated_node

    def leave_Assign(
        self,
        original_node: Assign,  # noqa: ARG002
        updated_node: Assign,
    ) -> BaseSmallStatement | FlattenSentinel[BaseSmallStatement] | RemovalSentinel:
        return self.handle_assignment(updated_node)

    def leave_SimpleStatementLine(
        self,
        original_node: SimpleStatementLine,  # noqa: ARG002
        updated_node: SimpleStatementLine,
    ) -> Any:  # noqa: ANN401
        return self.statement_handler.handle_statement(updated_node)

    def leave_SimpleStatementSuite(
        self,
        original_node: SimpleStatementSuite,  # noqa: ARG002
        updated_node: SimpleStatementSuite,
    ) -> Any:  # noqa: ANN401
        return self.statement_handler.handle_statement(updated_node)

    def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> RemovalSentinel | FunctionDef:
        body_block = updated_node.body

        # For IndentedBlock (regular multi-line functions)
        if hasattr(body_block, "body") and isinstance(body_block.body, (list, tuple)):
            body = body_block.body
            if not body:
                return RemoveFromParent()

            # Only remove test functions that only contain pass (useless test stubs)
            # Non-test functions with pass are legitimate and should be kept
            if original_node.name.value.startswith("test_") and len(body) == 1:
                first = body[0]
                # Check if it's a SimpleStatementLine with just Pass()
                if isinstance(first, SimpleStatementLine) and len(first.body) == 1 and isinstance(first.body[0], Pass):
                    return RemoveFromParent()
                # Check if it's a SimpleStatementSuite (one-liner) with just Pass()
                if isinstance(first, Pass):
                    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 is_direct_function_call_assert(self, assert_node: Assert) -> bool:
        return (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)
        )

    def handle_comparison_assert(self, comparison: Comparison) -> Assign | None:
        for side in [comparison.left, comparison.comparisons[0].comparator]:
            if self._should_create_codeflash_assignment(cast("BaseSmallStatement", side)):
                return self.create_codeflash_assignment(side)
            if isinstance(side, Name) and side.value in self.codeflash_references:
                return None
        return None

    def _should_create_codeflash_assignment(self, node: BaseSmallStatement) -> bool:
        return (
            self.contains_codeflash_reference(node)
            or (isinstance(node, ListComp) and isinstance(node.elt, Call) and self.is_target_function_node(node.elt))
            or (isinstance(node, Call) and self.is_target_function_node(node))
        )

    def create_codeflash_assignment(self, value_node: Any) -> Assign:  # noqa: ANN401
        return Assign(targets=[AssignTarget(target=Name("codeflash_output"))], value=value_node)

    def contains_codeflash_reference(self, node: Any) -> bool:  # noqa: ANN401, PLR0911
        if node is None:
            return False

        if isinstance(node, Name):
            return node.value == "codeflash_output" and node.value not in self.codeflash_references

        if isinstance(node, Call):
            return self.is_target_function_node(node) or any(
                self.contains_codeflash_reference(arg.value) for arg in node.args
            )

        if isinstance(node, ListComp):
            return isinstance(node.elt, Call) and self.is_target_function_node(node.elt)

        if isinstance(node, Attribute | Comparison | UnaryOperation):
            if isinstance(node, Attribute):
                return self.contains_codeflash_reference(node.value)
            if isinstance(node, Comparison):
                return self.contains_codeflash_reference(node.left) or any(
                    self.contains_codeflash_reference(comp.comparator) for comp in node.comparisons
                )
            return self.contains_codeflash_reference(node.expression)

        return False

    def is_variable_referenced(self, var_name: str) -> bool:
        return self.variable_reassignments.get(var_name) is not None or any(
            target == var_name for target in self.codeflash_references.values()
        )

    def leave_Module(self, original_node: Module, updated_node: Module) -> Module:  # noqa: ARG002
        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."
            ),
        )
        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)]
    result = modified_tree.with_changes(body=tuple(new_body))
    # Re-parse to ensure valid CST structure. The transformer may create nodes with
    # semicolon=None instead of MaybeSentinel, which causes errors during subsequent visits.
    return parse_module(result.code)
