"""Tests for graph representation and manipulation."""

from __future__ import annotations

import pytest

from sagellm_backend.graph.graph import Graph, Node, OpType


class TestNode:
    """Tests for Node class."""

    def test_node_creation(self) -> None:
        """Test basic node creation."""
        node = Node(node_id="n1", op_type=OpType.MATMUL)
        assert node.node_id == "n1"
        assert node.op_type == OpType.MATMUL
        assert node.inputs == []
        assert node.outputs == []
        assert node.device == "cpu"
        assert node.dtype == "float32"

    def test_node_with_connections(self) -> None:
        """Test node with input/output connections."""
        node = Node(
            node_id="n2",
            op_type=OpType.ADD,
            inputs=["n0", "n1"],
            outputs=["n3"],
        )
        assert len(node.inputs) == 2
        assert len(node.outputs) == 1

    def test_add_input(self) -> None:
        """Test adding input connections."""
        node = Node(node_id="n1", op_type=OpType.MATMUL)
        node.add_input("n0")
        assert "n0" in node.inputs

        # Adding same input again should be idempotent
        node.add_input("n0")
        assert node.inputs.count("n0") == 1

    def test_add_output(self) -> None:
        """Test adding output connections."""
        node = Node(node_id="n1", op_type=OpType.MATMUL)
        node.add_output("n2")
        assert "n2" in node.outputs

    def test_remove_input(self) -> None:
        """Test removing input connections."""
        node = Node(node_id="n1", op_type=OpType.ADD, inputs=["n0", "n1"])
        node.remove_input("n0")
        assert "n0" not in node.inputs
        assert "n1" in node.inputs

    def test_remove_output(self) -> None:
        """Test removing output connections."""
        node = Node(node_id="n1", op_type=OpType.ADD, outputs=["n2", "n3"])
        node.remove_output("n2")
        assert "n2" not in node.outputs
        assert "n3" in node.outputs

    def test_node_with_attrs(self) -> None:
        """Test node with operation-specific attributes."""
        node = Node(
            node_id="n1",
            op_type=OpType.MATMUL,
            attrs={"transpose_a": True, "transpose_b": False},
        )
        assert node.attrs["transpose_a"] is True
        assert node.attrs["transpose_b"] is False

    def test_node_with_metadata(self) -> None:
        """Test node with metadata."""
        node = Node(
            node_id="n1",
            op_type=OpType.ATTENTION,
            metadata={"num_heads": 8, "head_dim": 64},
        )
        assert node.metadata["num_heads"] == 8


class TestGraph:
    """Tests for Graph class."""

    def test_empty_graph(self) -> None:
        """Test creating an empty graph."""
        graph = Graph()
        assert len(graph.nodes) == 0
        assert len(graph.inputs) == 0
        assert len(graph.outputs) == 0

    def test_add_node(self) -> None:
        """Test adding nodes to graph."""
        graph = Graph()
        node = Node(node_id="n1", op_type=OpType.INPUT)
        graph.add_node(node)

        assert "n1" in graph.nodes
        assert graph.get_node("n1") == node

    def test_add_duplicate_node(self) -> None:
        """Test that adding duplicate node IDs raises error."""
        graph = Graph()
        node1 = Node(node_id="n1", op_type=OpType.INPUT)
        node2 = Node(node_id="n1", op_type=OpType.OUTPUT)

        graph.add_node(node1)
        with pytest.raises(ValueError, match="already exists"):
            graph.add_node(node2)

    def test_remove_node(self) -> None:
        """Test removing a node from graph."""
        graph = Graph()
        node = Node(node_id="n1", op_type=OpType.CONSTANT)
        graph.add_node(node)

        assert "n1" in graph.nodes
        graph.remove_node("n1")
        assert "n1" not in graph.nodes

    def test_add_edge(self) -> None:
        """Test adding edges between nodes."""
        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")

        assert "n2" in graph.nodes["n1"].outputs
        assert "n1" in graph.nodes["n2"].inputs

    def test_add_edge_missing_node(self) -> None:
        """Test that adding edge with missing node raises error."""
        graph = Graph()
        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        graph.add_node(n1)

        with pytest.raises(ValueError, match="not found"):
            graph.add_edge("n1", "n999")

    def test_remove_edge(self) -> None:
        """Test removing edges between nodes."""
        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")

        graph.remove_edge("n1", "n2")

        assert "n2" not in graph.nodes["n1"].outputs
        assert "n1" not in graph.nodes["n2"].inputs

    def test_topological_sort(self) -> None:
        """Test topological sorting of graph nodes."""
        graph = Graph()

        # Create a simple chain: n1 -> n2 -> n3
        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.ADD)
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_node(n3)
        graph.add_edge("n1", "n2")
        graph.add_edge("n2", "n3")

        order = graph.topological_sort()

        # n1 must come before n2, n2 before n3
        assert order.index("n1") < order.index("n2")
        assert order.index("n2") < order.index("n3")

    def test_topological_sort_with_cycle(self) -> None:
        """Test that topological sort detects cycles."""
        graph = Graph()

        # Create a cycle: n1 -> n2 -> n3 -> n1
        n1 = Node(node_id="n1", op_type=OpType.ADD)
        n2 = Node(node_id="n2", op_type=OpType.MUL)
        n3 = Node(node_id="n3", op_type=OpType.ADD)

        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_node(n3)
        graph.add_edge("n1", "n2")
        graph.add_edge("n2", "n3")
        graph.add_edge("n3", "n1")  # Creates cycle

        with pytest.raises(ValueError, match="cycle"):
            graph.topological_sort()

    def test_validate_valid_graph(self) -> None:
        """Test validation of a valid graph."""
        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.inputs = ["n1"]
        graph.outputs = ["n2"]

        assert graph.validate() is True

    def test_validate_missing_input(self) -> None:
        """Test validation fails with missing input node."""
        graph = Graph()
        graph.inputs = ["n999"]  # Non-existent node

        with pytest.raises(ValueError, match="Input node not found"):
            graph.validate()

    def test_validate_missing_output(self) -> None:
        """Test validation fails with missing output node."""
        graph = Graph()
        graph.outputs = ["n999"]  # Non-existent node

        with pytest.raises(ValueError, match="Output node not found"):
            graph.validate()

    def test_clone(self) -> None:
        """Test graph cloning."""
        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.inputs = ["n1"]
        graph.outputs = ["n2"]
        graph.metadata["model"] = "test"

        cloned = graph.clone()

        # Check structure is copied
        assert len(cloned.nodes) == len(graph.nodes)
        assert cloned.inputs == graph.inputs
        assert cloned.outputs == graph.outputs
        assert cloned.metadata == graph.metadata

        # Check it's a deep copy (modifying clone doesn't affect original)
        cloned.remove_node("n1")
        assert "n1" in graph.nodes
        assert "n1" not in cloned.nodes

    def test_complex_graph(self) -> None:
        """Test a more complex graph structure."""
        graph = Graph()

        # Create a diamond pattern:
        #     n1
        #    /  \\
        #   n2  n3
        #    \\  /
        #     n4
        nodes = [
            Node(node_id="n1", op_type=OpType.INPUT),
            Node(node_id="n2", op_type=OpType.MATMUL),
            Node(node_id="n3", op_type=OpType.ADD),
            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.inputs = ["n1"]
        graph.outputs = ["n4"]

        # Should validate successfully
        assert graph.validate() is True

        # Check topological order
        order = graph.topological_sort()
        assert order.index("n1") < order.index("n2")
        assert order.index("n1") < order.index("n3")
        assert order.index("n2") < order.index("n4")
        assert order.index("n3") < order.index("n4")
