"""Unit tests for DefaultIRExecutor."""

from __future__ import annotations

from types import SimpleNamespace

import pytest

from sagellm_control.ir.executor import DefaultIRExecutor, ExecutionCommand
from sagellm_control.ir.types import (
    ComputationNode,
    DependencyEdge,
    NodeType,
    ParallelismConfig,
    ParallelismStrategy,
    ResourceAllocation,
    SchedulerIR,
    TaskNode,
)


def _build_single_task_ir() -> SchedulerIR:
    ir = SchedulerIR(ir_id="ir-exec-1")
    task = TaskNode(
        node_id="task-req-1",
        node_type=NodeType.TASK,
        request_id="req-1",
        trace_id="trace-1",
        model_id="Qwen2-7B",
        prompt="hello executor",
        max_tokens=16,
    )
    prefill = ComputationNode(
        node_id="prefill-req-1",
        node_type=NodeType.PREFILL,
        task_id="task-req-1",
        operation="prefill",
        input_token_count=2,
        resource_allocation=ResourceAllocation(
            engine_id="engine-1",
            kv_cache_tokens=2,
            reuse_kv_prefix=True,
            kv_prefix_length=1,
        ),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )
    decode = ComputationNode(
        node_id="decode-req-1",
        node_type=NodeType.DECODE,
        task_id="task-req-1",
        operation="decode",
        output_token_count=16,
        resource_allocation=ResourceAllocation(engine_id="engine-1", kv_cache_tokens=16),
        parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2),
    )

    ir.add_node(task)
    ir.add_node(prefill)
    ir.add_node(decode)
    ir.add_edge(DependencyEdge(from_node=task.node_id, to_node=prefill.node_id))
    ir.add_edge(DependencyEdge(from_node=prefill.node_id, to_node=decode.node_id))
    return ir


class _MockControlPlane:
    async def execute_request(self, request: object) -> object:
        return SimpleNamespace(
            request_id=request.request_id,
            output_text="ok-control-plane",
            finish_reason="stop",
        )


class _MockEngineClient:
    async def execute(self, request: object) -> object:
        return SimpleNamespace(
            request_id=request.request_id,
            output_text="ok-engine",
            finish_reason="stop",
        )


class _FailingControlPlane:
    async def execute_request(self, request: object) -> object:
        raise RuntimeError("backend failure")


class TestExecutionCommand:
    """Tests for ExecutionCommand type."""

    def test_to_dict(self) -> None:
        """Should serialize command into dictionary."""
        command = ExecutionCommand(
            command_id="cmd-1",
            node_id="task-1",
            command_type="execute",
            engine_id="engine-1",
            payload={"k": "v"},
            metadata={"m": 1},
        )
        d = command.to_dict()
        assert d["command_id"] == "cmd-1"
        assert d["payload"]["k"] == "v"
        assert d["metadata"]["m"] == 1


class TestDefaultIRExecutor:
    """Tests for DefaultIRExecutor behavior."""

    def test_init_requires_backend(self) -> None:
        """Should require control_plane or engine_client."""
        with pytest.raises(ValueError):
            DefaultIRExecutor()

    def test_translate_ir_to_commands(self) -> None:
        """Should translate task IR into a single execute command."""
        ir = _build_single_task_ir()
        executor = DefaultIRExecutor(control_plane=_MockControlPlane())

        commands = executor.translate_to_commands(ir)
        assert len(commands) == 1
        command = commands[0]
        assert command.command_id == "cmd-req-1"
        assert command.command_type == "execute"
        assert command.payload["model"] == "Qwen2-7B"
        assert command.metadata["reuse_kv_prefix"] is True
        assert command.metadata["kv_prefix_length"] == 1

    @pytest.mark.asyncio
    async def test_execute_via_control_plane(self) -> None:
        """Should execute commands through control plane backend."""
        ir = _build_single_task_ir()
        executor = DefaultIRExecutor(control_plane=_MockControlPlane())

        result = await executor.execute(ir)
        assert "cmd-req-1" in result
        assert result["cmd-req-1"]["output_text"] == "ok-control-plane"

    @pytest.mark.asyncio
    async def test_execute_via_engine_client(self) -> None:
        """Should execute commands through direct engine client backend."""
        ir = _build_single_task_ir()
        executor = DefaultIRExecutor(engine_client=_MockEngineClient())

        result = await executor.execute(ir)
        assert "cmd-req-1" in result
        assert result["cmd-req-1"]["output_text"] == "ok-engine"

    @pytest.mark.asyncio
    async def test_batch_execute_multiple_commands(self) -> None:
        """Should execute all commands produced from multiple tasks."""
        ir = _build_single_task_ir()
        second_task = TaskNode(
            node_id="task-req-2",
            node_type=NodeType.TASK,
            request_id="req-2",
            trace_id="trace-2",
            model_id="Qwen2-7B",
            prompt="second prompt",
            max_tokens=8,
        )
        second_prefill = ComputationNode(
            node_id="prefill-req-2",
            node_type=NodeType.PREFILL,
            task_id="task-req-2",
            operation="prefill",
            input_token_count=2,
            resource_allocation=ResourceAllocation(engine_id="engine-2", kv_cache_tokens=2),
            parallelism=ParallelismConfig(
                strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2
            ),
        )
        second_decode = ComputationNode(
            node_id="decode-req-2",
            node_type=NodeType.DECODE,
            task_id="task-req-2",
            operation="decode",
            output_token_count=8,
            resource_allocation=ResourceAllocation(engine_id="engine-2", kv_cache_tokens=8),
            parallelism=ParallelismConfig(
                strategy=ParallelismStrategy.TENSOR_PARALLEL, tp_degree=2
            ),
        )
        ir.add_node(second_task)
        ir.add_node(second_prefill)
        ir.add_node(second_decode)
        ir.add_edge(DependencyEdge(from_node=second_task.node_id, to_node=second_prefill.node_id))
        ir.add_edge(DependencyEdge(from_node=second_prefill.node_id, to_node=second_decode.node_id))

        executor = DefaultIRExecutor(control_plane=_MockControlPlane())
        result = await executor.execute(ir)

        assert "cmd-req-1" in result
        assert "cmd-req-2" in result
        assert len(result) == 2

    @pytest.mark.asyncio
    async def test_execute_error_handling(self) -> None:
        """Should return error result when command execution fails."""
        ir = _build_single_task_ir()
        executor = DefaultIRExecutor(control_plane=_FailingControlPlane())

        result = await executor.execute(ir)
        assert "cmd-req-1" in result
        assert "error" in result["cmd-req-1"]
        assert "backend failure" in result["cmd-req-1"]["error"]
