"""Tests for CPU-specific optimization passes."""

from __future__ import annotations

from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.passes.constant_folding import ConstantFoldingPass
from sagellm_backend.graph.passes.dead_node_elimination import DeadNodeEliminationPass
from sagellm_backend.graph.passes.identity_elimination import IdentityEliminationPass


class TestDeadNodeEliminationPass:
    """Tests for DeadNodeEliminationPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = DeadNodeEliminationPass()
        assert pass_obj.name == "dead_node_elimination"

    def test_no_dead_nodes(self) -> None:
        """Test graph with no dead nodes."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.OUTPUT)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")
        graph.outputs = ["n2"]

        result = pass_obj.run(graph)

        assert result.modified is False
        assert len(graph.nodes) == 2

    def test_remove_single_dead_node(self) -> None:
        """Test removing a single dead node."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.OUTPUT)
        n3 = Node(node_id="n3", op_type=OpType.ADD)  # Dead - not connected to output

        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_node(n3)
        graph.add_edge("n1", "n2")
        graph.outputs = ["n2"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1
        assert len(graph.nodes) == 2
        assert "n3" not in graph.nodes

    def test_remove_dead_subgraph(self) -> None:
        """Test removing a dead subgraph."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        # Live path: n1 -> n2 -> n3
        # Dead path: n4 -> n5 -> n6
        live_nodes = [
            Node(node_id="n1", op_type=OpType.INPUT),
            Node(node_id="n2", op_type=OpType.ADD),
            Node(node_id="n3", op_type=OpType.OUTPUT),
        ]
        dead_nodes = [
            Node(node_id="n4", op_type=OpType.INPUT),
            Node(node_id="n5", op_type=OpType.MUL),
            Node(node_id="n6", op_type=OpType.OUTPUT),
        ]

        for node in live_nodes + dead_nodes:
            graph.add_node(node)

        graph.add_edge("n1", "n2")
        graph.add_edge("n2", "n3")
        graph.add_edge("n4", "n5")
        graph.add_edge("n5", "n6")

        # Only n3 is marked as output
        graph.outputs = ["n3"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 3  # n4, n5, n6 removed
        assert len(graph.nodes) == 3
        assert "n4" not in graph.nodes
        assert "n5" not in graph.nodes
        assert "n6" not in graph.nodes

    def test_keep_all_paths_to_outputs(self) -> None:
        """Test that all paths to outputs are kept."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        # Diamond pattern: n1 splits to n2 and n3, both feed to n4
        #     n1
        #    /  \\
        #   n2  n3
        #    \\  /
        #     n4
        nodes = [
            Node(node_id="n1", op_type=OpType.INPUT),
            Node(node_id="n2", op_type=OpType.ADD),
            Node(node_id="n3", op_type=OpType.MUL),
            Node(node_id="n4", op_type=OpType.OUTPUT),
        ]

        for node in nodes:
            graph.add_node(node)

        graph.add_edge("n1", "n2")
        graph.add_edge("n1", "n3")
        graph.add_edge("n2", "n4")
        graph.add_edge("n3", "n4")
        graph.outputs = ["n4"]

        result = pass_obj.run(graph)

        # No nodes should be removed
        assert result.modified is False
        assert len(graph.nodes) == 4


class TestConstantFoldingPass:
    """Tests for Co τnstantFoldingPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = ConstantFoldingPass()
        assert pass_obj.name == "constant_folding"

    def test_no_foldable_operations(self) -> None:
        """Test graph with no foldable operations."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.ADD)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")

        result = pass_obj.run(graph)

        assert result.modified is False

    def test_fold_constant_add(self) -> None:
        """Test folding constant addition."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        # Two constants feeding into an ADD
        c1 = Node(node_id="c1", op_type=OpType.CONSTANT)
        c2 = Node(node_id="c2", op_type=OpType.CONSTANT)
        add = Node(node_id="add1", op_type=OpType.ADD)

        graph.add_node(c1)
        graph.add_node(c2)
        graph.add_node(add)
        graph.add_edge("c1", "add1")
        graph.add_edge("c2", "add1")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1

        # ADD should be converted to CONSTANT
        assert graph.nodes["add1"].op_type == OpType.CONSTANT
        assert graph.nodes["add1"].metadata["folded"] is True

    def test_fold_multiple_operations(self) -> None:
        """Test folding multiple constant operations."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        # Multiple constant operations
        c1 = Node(node_id="c1", op_type=OpType.CONSTANT)
        c2 = Node(node_id="c2", op_type=OpType.CONSTANT)
        c3 = Node(node_id="c3", op_type=OpType.CONSTANT)
        add1 = Node(node_id="add1", op_type=OpType.ADD)
        mul1 = Node(node_id="mul1", op_type=OpType.MUL)

        graph.add_node(c1)
        graph.add_node(c2)
        graph.add_node(c3)
        graph.add_node(add1)
        graph.add_node(mul1)

        graph.add_edge("c1", "add1")
        graph.add_edge("c2", "add1")
        graph.add_edge("c3", "mul1")
        graph.add_edge("add1", "mul1")

        # Run once - should fold add1
        result1 = pass_obj.run(graph)
        assert result1.modified is True

        # Run again - should fold mul1 (now both inputs are constant)
        result2 = pass_obj.run(graph)
        assert result2.modified is True

    def test_skip_already_constant(self) -> None:
        """Test that constant nodes are not folded again."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        c1 = Node(node_id="c1", op_type=OpType.CONSTANT)
        graph.add_node(c1)

        result = pass_obj.run(graph)

        assert result.modified is False


class TestIdentityEliminationPass:
    """Tests for IdentityEliminationPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = IdentityEliminationPass()
        assert pass_obj.name == "identity_elimination"

    def test_no_identity_operations(self) -> None:
        """Test graph with no identity operations."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.ADD)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")

        result = pass_obj.run(graph)

        assert result.modified is False

    def test_eliminate_identity_reshape(self) -> None:
        """Test eliminating identity reshape."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        reshape = Node(
            node_id="reshape1",
            op_type=OpType.RESHAPE,
            attrs={"is_identity": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(reshape)
        graph.add_node(n3)
        graph.add_edge("n1", "reshape1")
        graph.add_edge("reshape1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1
        assert "reshape1" not in graph.nodes

        # n1 should connect directly to n3
        assert "n3" in graph.nodes["n1"].outputs
        assert "n1" in graph.nodes["n3"].inputs

    def test_eliminate_identity_transpose(self) -> None:
        """Test eliminating identity transpose."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        transpose = Node(
            node_id="transpose1",
            op_type=OpType.TRANSPOSE,
            attrs={"perm": [0, 1, 2]},  # Identity permutation
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(transpose)
        graph.add_node(n3)
        graph.add_edge("n1", "transpose1")
        graph.add_edge("transpose1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "transpose1" not in graph.nodes

    def test_eliminate_add_zero(self) -> None:
        """Test eliminating ADD with zero."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        add = Node(
            node_id="add1",
            op_type=OpType.ADD,
            attrs={"add_zero": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(add)
        graph.add_node(n3)
        graph.add_edge("n1", "add1")
        graph.add_edge("add1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "add1" not in graph.nodes

    def test_eliminate_mul_one(self) -> None:
        """Test eliminating MUL with one."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        mul = Node(
            node_id="mul1",
            op_type=OpType.MUL,
            attrs={"mul_one": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(mul)
        graph.add_node(n3)
        graph.add_edge("n1", "mul1")
        graph.add_edge("mul1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "mul1" not in graph.nodes

    def test_eliminate_chain_of_identities(self) -> None:
        """Test eliminating multiple identity operations in sequence."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        id1 = Node(node_id="id1", op_type=OpType.RESHAPE, attrs={"is_identity": True})
        id2 = Node(node_id="id2", op_type=OpType.ADD, attrs={"add_zero": True})
        n4 = Node(node_id="n4", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(id1)
        graph.add_node(id2)
        graph.add_node(n4)

        graph.add_edge("n1", "id1")
        graph.add_edge("id1", "id2")
        graph.add_edge("id2", "n4")

        # Run multiple times to eliminate chain
        result1 = pass_obj.run(graph)
        assert result1.modified is True
        assert result1.num_changes >= 1

        _ = pass_obj.run(graph)
        # May eliminate more or be done
        assert len(graph.nodes) <= 4
