"""Unit tests for IROptimizer and optimization passes."""

from __future__ import annotations

from sagellm_control.ir.optimizer import (
    ComputeCommOverlapPass,
    IROptimizer,
    KVReusePass,
    OptimizationPass,
    ResourceAllocationPass,
)
from sagellm_control.ir.types import (
    ComputationNode,
    DependencyEdge,
    NodeType,
    ParallelismConfig,
    ParallelismStrategy,
    ResourceAllocation,
    SchedulerIR,
    TaskNode,
)


def _build_two_task_ir() -> SchedulerIR:
    ir = SchedulerIR(ir_id="ir-opt-1")

    task_a = TaskNode(
        node_id="task-a",
        node_type=NodeType.TASK,
        request_id="req-a",
        trace_id="trace-a",
        model_id="Qwen2-7B",
        prompt="shared prefix tokens alpha beta",
        max_tokens=16,
    )
    task_b = TaskNode(
        node_id="task-b",
        node_type=NodeType.TASK,
        request_id="req-b",
        trace_id="trace-b",
        model_id="Qwen2-7B",
        prompt="shared prefix tokens gamma delta",
        max_tokens=16,
    )
    prefill_a = ComputationNode(
        node_id="prefill-a",
        node_type=NodeType.PREFILL,
        task_id="task-a",
        operation="prefill",
        input_token_count=5,
        resource_allocation=ResourceAllocation(kv_cache_tokens=70, memory_mb=100.0),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )
    decode_a = ComputationNode(
        node_id="decode-a",
        node_type=NodeType.DECODE,
        task_id="task-a",
        operation="decode",
        output_token_count=16,
        resource_allocation=ResourceAllocation(kv_cache_tokens=30, memory_mb=50.0),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )
    prefill_b = ComputationNode(
        node_id="prefill-b",
        node_type=NodeType.PREFILL,
        task_id="task-b",
        operation="prefill",
        input_token_count=5,
        resource_allocation=ResourceAllocation(kv_cache_tokens=70, memory_mb=100.0),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )
    decode_b = ComputationNode(
        node_id="decode-b",
        node_type=NodeType.DECODE,
        task_id="task-b",
        operation="decode",
        output_token_count=16,
        resource_allocation=ResourceAllocation(kv_cache_tokens=30, memory_mb=50.0),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )

    for node in (task_a, prefill_a, decode_a, task_b, prefill_b, decode_b):
        ir.add_node(node)

    ir.add_edge(DependencyEdge(from_node="task-a", to_node="prefill-a"))
    ir.add_edge(DependencyEdge(from_node="prefill-a", to_node="decode-a"))
    ir.add_edge(DependencyEdge(from_node="task-b", to_node="prefill-b"))
    ir.add_edge(DependencyEdge(from_node="prefill-b", to_node="decode-b"))

    ir.global_parallelism = ParallelismConfig(
        strategy=ParallelismStrategy.TENSOR_PARALLEL,
        tp_degree=2,
        pp_degree=1,
        dp_degree=1,
    )
    return ir


class _MarkerPass(OptimizationPass):
    """Simple pass for pipeline tests."""

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

    def optimize(self, ir: SchedulerIR) -> SchedulerIR:
        ir.optimization_hints["marker_ran"] = True
        return ir


class _FailingPass(OptimizationPass):
    """Pass that fails intentionally for resiliency test."""

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

    def optimize(self, ir: SchedulerIR) -> SchedulerIR:
        raise RuntimeError("boom")


class TestKVReusePass:
    """Tests for KVReusePass."""

    def test_detect_common_prefix_marks_prefill_reuse(self) -> None:
        """Should mark KV reuse when common prefix exceeds threshold."""
        ir = _build_two_task_ir()
        pass_ = KVReusePass(min_reuse_length=2)

        optimized = pass_.optimize(ir)
        prefill_b = optimized.get_node("prefill-b")
        assert prefill_b is not None
        assert prefill_b.resource_allocation.reuse_kv_prefix is True
        assert prefill_b.resource_allocation.kv_prefix_length >= 2

    def test_no_reuse_when_prefix_too_short(self) -> None:
        """Should not mark reuse if threshold is too high."""
        ir = _build_two_task_ir()
        pass_ = KVReusePass(min_reuse_length=999)

        optimized = pass_.optimize(ir)
        prefill_b = optimized.get_node("prefill-b")
        assert prefill_b is not None
        assert prefill_b.resource_allocation.reuse_kv_prefix is False


class TestComputeCommOverlapPass:
    """Tests for ComputeCommOverlapPass."""

    def test_add_overlap_hints_when_parallelism_enabled(self) -> None:
        """Should annotate computation nodes for overlap."""
        ir = _build_two_task_ir()
        pass_ = ComputeCommOverlapPass()

        optimized = pass_.optimize(ir)
        for node in optimized.get_computation_nodes():
            assert node.metadata.get("enable_tp_overlap") is True


class TestResourceAllocationPass:
    """Tests for ResourceAllocationPass."""

    def test_reduce_kv_allocation_when_over_budget(self) -> None:
        """Should scale down KV allocation to fit budget."""
        ir = _build_two_task_ir()
        pass_ = ResourceAllocationPass(kv_cache_budget=100, memory_budget_mb=10_000.0)

        optimized = pass_.optimize(ir)
        total_kv = sum(
            node.resource_allocation.kv_cache_tokens for node in optimized.get_computation_nodes()
        )
        assert total_kv <= 100


class TestIROptimizer:
    """Tests for IROptimizer pipeline and pass management."""

    def test_optimizer_pipeline_execution(self) -> None:
        """Should execute passes in pipeline."""
        ir = _build_two_task_ir()
        optimizer = IROptimizer(passes=[_MarkerPass()], enable_default_passes=False)

        optimized = optimizer.optimize(ir)
        assert optimized.optimization_hints["marker_ran"] is True

    def test_optimizer_continues_after_pass_failure(self) -> None:
        """Should continue running later passes when one fails."""
        ir = _build_two_task_ir()
        optimizer = IROptimizer(
            passes=[_FailingPass(), _MarkerPass()],
            enable_default_passes=False,
        )

        optimized = optimizer.optimize(ir)
        assert optimized.optimization_hints["marker_ran"] is True

    def test_add_remove_and_get_pass(self) -> None:
        """Should support pass registration lifecycle."""
        optimizer = IROptimizer(enable_default_passes=False)
        marker = _MarkerPass()

        optimizer.add_pass(marker)
        assert optimizer.get_pass("marker") is marker
        assert "marker" in optimizer.list_passes()
        assert optimizer.remove_pass("marker") is True
        assert optimizer.get_pass("marker") is None
        assert optimizer.remove_pass("marker") is False
