"""Tests for graph optimizer."""

from __future__ import annotations

import pytest

from sagellm_backend.graph.config import OptimizationConfig
from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.optimizer import GraphOptimizer, OptimizationReport
from sagellm_backend.graph.pass_base import (
    OptimizationPass,
    PassRegistry,
    PassResult,
    register_pass,
)


class TestOptimizationReport:
    """Tests for OptimizationReport."""

    def test_default_report(self) -> None:
        """Test default report."""
        report = OptimizationReport()
        assert report.total_time_ms == 0.0
        assert report.num_iterations == 0
        assert report.num_passes == 0
        assert report.modifications == 0


class TestGraphOptimizer:
    """Tests for GraphOptimizer."""

    def setup_method(self) -> None:
        """Clear pass registry before each test."""
        PassRegistry.clear()

    def test_optimizer_initialization(self) -> None:
        """Test optimizer initialization."""
        config = OptimizationConfig.default_cpu()
        optimizer = GraphOptimizer(config)

        assert optimizer.config == config

    def test_optimizer_default_config(self) -> None:
        """Test optimizer with default config."""
        optimizer = GraphOptimizer()
        assert optimizer.config.level == 1
        assert optimizer.config.target_device == "cpu"

    def test_optimize_empty_graph(self) -> None:
        """Test optimizing an empty graph."""
        optimizer = GraphOptimizer()
        graph = Graph()

        optimized, report = optimizer.optimize(graph)

        assert len(optimized.nodes) == 0
        assert report.initial_num_nodes == 0
        assert report.final_num_nodes == 0

    def test_optimize_simple_graph(self) -> None:
        """Test optimizing a simple graph."""

        # Register a simple pass
        @register_pass
        class SimplePass(OptimizationPass):
            @property
            def name(self) -> str:
                return "simple_pass"

            def run(self, graph: Graph) -> PassResult:
                return PassResult(modified=False, message="No changes")

        optimizer = GraphOptimizer()
        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"]

        optimized, report = optimizer.optimize(graph)

        assert len(optimized.nodes) == 2
        assert report.num_passes > 0
        assert report.total_time_ms > 0

    def test_optimize_with_modifications(self) -> None:
        """Test optimization that modifies the graph."""

        @register_pass
        class ModifyingPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "modifying_pass"

            def run(self, graph: Graph) -> PassResult:
                # Add a metadata flag (only modify once)
                if any("optimized" in node.metadata for node in graph.nodes.values()):
                    return PassResult(modified=False, message="Already optimized")

                for node in graph.nodes.values():
                    node.metadata["optimized"] = True

                return PassResult(
                    modified=True,
                    num_changes=len(graph.nodes),
                    message=f"Modified {len(graph.nodes)} nodes",
                )

        optimizer = GraphOptimizer()
        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.inputs = ["n1"]
        graph.outputs = ["n2"]

        optimized, report = optimizer.optimize(graph)

        assert report.modifications == 2
        assert optimized.nodes["n1"].metadata["optimized"] is True
        assert optimized.nodes["n2"].metadata["optimized"] is True

    def test_optimize_convergence(self) -> None:
        """Test that optimizer stops when no modifications are made."""

        call_count = 0

        @register_pass
        class OnceModifyingPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "once_modifying_pass"

            def run(self, graph: Graph) -> PassResult:
                nonlocal call_count
                call_count += 1

                # Only modify on first call
                if call_count == 1:
                    return PassResult(modified=True, num_changes=1)
                return PassResult(modified=False)

        config = OptimizationConfig(max_iterations=10)
        optimizer = GraphOptimizer(config)
        graph = Graph()

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

        optimized, report = optimizer.optimize(graph)

        # Should converge after 2 iterations (1st modifies, 2nd doesn't)
        assert report.num_iterations == 2
        assert call_count == 2

    def test_optimize_max_iterations(self) -> None:
        """Test that optimizer respects max_iterations."""

        @register_pass
        class AlwaysModifyingPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "always_modifying_pass"

            def run(self, graph: Graph) -> PassResult:
                # Always report modifications
                return PassResult(modified=True, num_changes=1)

        config = OptimizationConfig(max_iterations=3)
        optimizer = GraphOptimizer(config)
        graph = Graph()

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

        optimized, report = optimizer.optimize(graph)

        # Should stop at max_iterations
        assert report.num_iterations == 3

    def test_optimize_with_validation(self) -> None:
        """Test optimization with validation after each pass."""

        @register_pass
        class ValidatedPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "validated_pass"

            def run(self, graph: Graph) -> PassResult:
                return PassResult(modified=False)

        config = OptimizationConfig(validate_after_each_pass=True)
        optimizer = GraphOptimizer(config)
        graph = Graph()

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

        # Should not raise (validation passes)
        optimized, report = optimizer.optimize(graph)
        assert report.num_passes > 0

    def test_optimize_invalid_graph(self) -> None:
        """Test that optimizing invalid graph raises error."""
        optimizer = GraphOptimizer()
        graph = Graph()
        graph.inputs = ["n999"]  # Non-existent node

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

    def test_optimize_pass_produces_invalid_graph(self) -> None:
        """Test that pass producing invalid graph is caught."""

        @register_pass
        class InvalidatingPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "invalidating_pass"

            def run(self, graph: Graph) -> PassResult:
                # Create an invalid edge (reference non-existent node)
                # This will make the graph invalid
                for node in graph.nodes.values():
                    node.inputs.append("nonexistent_node_999")
                    break
                return PassResult(modified=True, num_changes=1)

        config = OptimizationConfig(validate_after_each_pass=True)
        optimizer = GraphOptimizer(config)
        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"]

        with pytest.raises(RuntimeError, match="(produced invalid graph|failed)"):
            optimizer.optimize(graph)

    def test_get_enabled_passes(self) -> None:
        """Test getting enabled pass names."""

        @register_pass
        class Pass1(OptimizationPass):
            @property
            def name(self) -> str:
                return "pass1"

            def run(self, graph: Graph) -> PassResult:
                return PassResult(modified=False)

        @register_pass
        class Pass2(OptimizationPass):
            @property
            def name(self) -> str:
                return "pass2"

            def run(self, graph: Graph) -> PassResult:
                return PassResult(modified=False)

        config = OptimizationConfig(enabled_passes=["pass1"])
        optimizer = GraphOptimizer(config)

        enabled = optimizer.get_enabled_passes()
        assert "pass1" in enabled
        assert "pass2" not in enabled

    def test_dependency_resolution(self) -> None:
        """Test that pass dependencies are resolved correctly."""

        @register_pass
        class BasePass(OptimizationPass):
            @property
            def name(self) -> str:
                return "base_pass"

            def run(self, graph: Graph) -> PassResult:
                graph.metadata["base_executed"] = True
                return PassResult(modified=True, num_changes=1)

        @register_pass
        class DependentPass(OptimizationPass):
            @property
            def name(self) -> str:
                return "dependent_pass"

            @property
            def requires(self) -> list[str]:
                return ["base_pass"]

            def run(self, graph: Graph) -> PassResult:
                # Check that base_pass ran first
                assert graph.metadata.get("base_executed") is True
                return PassResult(modified=False)

        config = OptimizationConfig(enabled_passes=["dependent_pass"])
        optimizer = GraphOptimizer(config)
        graph = Graph()

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

        optimized, report = optimizer.optimize(graph)

        # Both passes should have executed
        enabled = optimizer.get_enabled_passes()
        assert "base_pass" in enabled
        assert "dependent_pass" in enabled
