"""Tests for optimization pass infrastructure."""

from __future__ import annotations

import pytest

from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.pass_base import (
    OptimizationPass,
    PassRegistry,
    PassResult,
    register_pass,
)


class DummyPass(OptimizationPass):
    """A dummy pass for testing."""

    @property
    def name(self) -> str:
        return "dummy_pass"

    def run(self, graph: Graph) -> PassResult:
        return PassResult(modified=False, message="Dummy pass executed")


class ModifyingPass(OptimizationPass):
    """A pass that modifies the graph."""

    @property
    def name(self) -> str:
        return "modifying_pass"

    def run(self, graph: Graph) -> PassResult:
        # Add a metadata tag to all nodes
        for node in graph.nodes.values():
            node.metadata["modified"] = True
        return PassResult(
            modified=True,
            num_changes=len(graph.nodes),
            message=f"Modified {len(graph.nodes)} nodes",
        )


class DependentPass(OptimizationPass):
    """A pass with dependencies."""

    @property
    def name(self) -> str:
        return "dependent_pass"

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

    def run(self, graph: Graph) -> PassResult:
        return PassResult(modified=False, message="Dependent pass executed")


class HighPriorityPass(OptimizationPass):
    """A high priority pass."""

    @property
    def name(self) -> str:
        return "high_priority_pass"

    @property
    def priority(self) -> int:
        return 100

    def run(self, graph: Graph) -> PassResult:
        return PassResult(modified=False, message="High priority pass")


class TestPassResult:
    """Tests for PassResult."""

    def test_default_result(self) -> None:
        """Test default PassResult."""
        result = PassResult()
        assert result.modified is False
        assert result.num_changes == 0
        assert result.metrics == {}
        assert result.message == ""

    def test_result_with_changes(self) -> None:
        """Test PassResult with changes."""
        result = PassResult(
            modified=True,
            num_changes=5,
            metrics={"nodes_removed": 3, "edges_added": 2},
            message="Optimized graph",
        )
        assert result.modified is True
        assert result.num_changes == 5
        assert result.metrics["nodes_removed"] == 3


class TestOptimizationPass:
    """Tests for OptimizationPass base class."""

    def test_pass_initialization(self) -> None:
        """Test pass initialization."""
        pass_obj = DummyPass()
        assert pass_obj.name == "dummy_pass"
        assert pass_obj.enabled_by_default is True
        assert pass_obj.requires == []
        assert pass_obj.priority == 0

    def test_pass_run(self) -> None:
        """Test running a pass."""
        pass_obj = DummyPass()
        graph = Graph()
        result = pass_obj.run(graph)

        assert isinstance(result, PassResult)
        assert result.message == "Dummy pass executed"

    def test_modifying_pass(self) -> None:
        """Test a pass that modifies the graph."""
        pass_obj = ModifyingPass()
        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)

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 2
        assert graph.nodes["n1"].metadata["modified"] is True
        assert graph.nodes["n2"].metadata["modified"] is True

    def test_pass_with_dependencies(self) -> None:
        """Test pass with dependencies."""
        pass_obj = DependentPass()
        assert pass_obj.requires == ["dummy_pass"]

    def test_pass_with_priority(self) -> None:
        """Test pass with custom priority."""
        pass_obj = HighPriorityPass()
        assert pass_obj.priority == 100


class TestPassRegistry:
    """Tests for PassRegistry."""

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

    def test_register_pass(self) -> None:
        """Test registering a pass."""
        PassRegistry.register(DummyPass)
        assert "dummy_pass" in PassRegistry.get_pass_names()

    def test_register_invalid_type(self) -> None:
        """Test that registering non-pass type raises error."""

        class NotAPass:
            pass

        with pytest.raises(TypeError, match="must inherit from OptimizationPass"):
            PassRegistry.register(NotAPass)  # type: ignore

    def test_get_pass(self) -> None:
        """Test retrieving a pass by name."""
        PassRegistry.register(DummyPass)
        pass_obj = PassRegistry.get_pass("dummy_pass")

        assert isinstance(pass_obj, DummyPass)
        assert pass_obj.name == "dummy_pass"

    def test_get_nonexistent_pass(self) -> None:
        """Test that getting non-existent pass raises error."""
        with pytest.raises(ValueError, match="Unknown pass"):
            PassRegistry.get_pass("nonexistent_pass")

    def test_get_all_passes(self) -> None:
        """Test getting all registered passes."""
        PassRegistry.register(DummyPass)
        PassRegistry.register(ModifyingPass)
        PassRegistry.register(HighPriorityPass)

        passes = PassRegistry.get_all_passes()

        assert len(passes) == 3
        # Should be sorted by priority (high priority first)
        assert passes[0].name == "high_priority_pass"

    def test_get_pass_names(self) -> None:
        """Test getting all pass names."""
        PassRegistry.register(DummyPass)
        PassRegistry.register(ModifyingPass)

        names = PassRegistry.get_pass_names()

        assert "dummy_pass" in names
        assert "modifying_pass" in names
        assert len(names) == 2

    def test_resolve_dependencies(self) -> None:
        """Test dependency resolution."""
        PassRegistry.register(DummyPass)
        PassRegistry.register(DependentPass)

        # Request only dependent_pass
        resolved = PassRegistry.resolve_dependencies(["dependent_pass"])

        # Should include dummy_pass (dependency)
        assert "dummy_pass" in resolved
        assert "dependent_pass" in resolved
        # dummy_pass should come before dependent_pass
        assert resolved.index("dummy_pass") < resolved.index("dependent_pass")

    def test_resolve_missing_dependency(self) -> None:
        """Test that missing dependency raises error."""
        PassRegistry.register(DependentPass)

        with pytest.raises(ValueError, match="Required pass not found"):
            PassRegistry.resolve_dependencies(["dependent_pass"])

    def test_clear_registry(self) -> None:
        """Test clearing the registry."""
        PassRegistry.register(DummyPass)
        assert len(PassRegistry.get_pass_names()) > 0

        PassRegistry.clear()
        assert len(PassRegistry.get_pass_names()) == 0

    def test_register_pass_decorator(self) -> None:
        """Test using @register_pass decorator."""

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

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

        # Should be automatically registered
        assert "decorated_pass" in PassRegistry.get_pass_names()
        pass_obj = PassRegistry.get_pass("decorated_pass")
        assert isinstance(pass_obj, DecoratedPass)
