"""Unit tests for Scheduler IR core types."""

from __future__ import annotations

from datetime import datetime

import pytest

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


class TestParallelismConfig:
    """Tests for ParallelismConfig."""

    def test_init_and_total_parallelism(self) -> None:
        """Should initialize config and compute total parallelism."""
        config = ParallelismConfig(
            strategy=ParallelismStrategy.HYBRID,
            tp_degree=2,
            pp_degree=2,
            dp_degree=2,
            ep_degree=1,
        )
        assert config.strategy == ParallelismStrategy.HYBRID
        assert config.total_parallelism == 8

    @pytest.mark.parametrize(
        ("field", "value"),
        [
            ("tp_degree", 0),
            ("pp_degree", 0),
            ("dp_degree", 0),
            ("ep_degree", 0),
        ],
    )
    def test_validation_fails_for_invalid_degree(self, field: str, value: int) -> None:
        """Should reject invalid parallelism degrees."""
        kwargs = {
            "tp_degree": 1,
            "pp_degree": 1,
            "dp_degree": 1,
            "ep_degree": 1,
        }
        kwargs[field] = value
        with pytest.raises(ValueError):
            ParallelismConfig(**kwargs)

    def test_to_dict_from_dict_roundtrip(self) -> None:
        """Should roundtrip via dictionary serialization."""
        config = ParallelismConfig(
            strategy=ParallelismStrategy.TENSOR_PARALLEL,
            tp_degree=4,
            pp_degree=1,
            dp_degree=2,
            ep_degree=1,
        )
        restored = ParallelismConfig.from_dict(config.to_dict())
        assert restored == config


class TestResourceAllocation:
    """Tests for ResourceAllocation."""

    def test_to_dict_from_dict_roundtrip(self) -> None:
        """Should serialize and deserialize ResourceAllocation."""
        allocation = ResourceAllocation(
            engine_id="engine-1",
            kv_cache_tokens=1024,
            memory_mb=512.5,
            compute_units=2,
            reuse_kv_prefix=True,
            kv_prefix_length=128,
        )
        restored = ResourceAllocation.from_dict(allocation.to_dict())
        assert restored == allocation


class TestNodeCreation:
    """Tests for IRNode, TaskNode, and ComputationNode."""

    def test_ir_node_create_and_serialize(self) -> None:
        """Should create and roundtrip base IRNode."""
        node = IRNode(node_id="n1", node_type=NodeType.SYNC, metadata={"k": "v"})
        restored = IRNode.from_dict(node.to_dict())
        assert restored.node_id == "n1"
        assert restored.node_type == NodeType.SYNC
        assert restored.metadata == {"k": "v"}

    def test_task_node_create(self) -> None:
        """Should create TaskNode and preserve fields."""
        node = TaskNode(
            node_id="task-1",
            node_type=NodeType.TASK,
            request_id="req-1",
            trace_id="trace-1",
            model_id="Qwen2-7B",
            prompt="hello world",
            max_tokens=64,
            priority=1,
            estimated_tokens=66,
        )
        assert node.request_id == "req-1"
        assert node.node_type == NodeType.TASK
        restored = TaskNode.from_dict(node.to_dict())
        assert restored.request_id == "req-1"
        assert restored.priority == 1

    def test_computation_node_create(self) -> None:
        """Should create ComputationNode and preserve nested fields."""
        node = ComputationNode(
            node_id="prefill-1",
            node_type=NodeType.PREFILL,
            task_id="task-1",
            operation="prefill",
            input_token_count=32,
            output_token_count=0,
            estimated_latency_ms=21.0,
            resource_allocation=ResourceAllocation(kv_cache_tokens=32),
            parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL),
        )
        restored = ComputationNode.from_dict(node.to_dict())
        assert restored.task_id == "task-1"
        assert restored.resource_allocation.kv_cache_tokens == 32
        assert restored.parallelism.strategy == ParallelismStrategy.TENSOR_PARALLEL


class TestDependencyEdge:
    """Tests for DependencyEdge."""

    def test_to_dict_from_dict_roundtrip(self) -> None:
        """Should serialize and deserialize edge."""
        edge = DependencyEdge(
            from_node="task-1",
            to_node="prefill-1",
            dependency_type="task_decomposition",
            weight=2.0,
        )
        restored = DependencyEdge.from_dict(edge.to_dict())
        assert restored == edge


class TestSchedulerIR:
    """Tests for SchedulerIR graph operations and serialization."""

    def _build_sample_ir(self) -> SchedulerIR:
        now = datetime.now()
        task = TaskNode(
            node_id="task-1",
            node_type=NodeType.TASK,
            created_at=now,
            request_id="req-1",
            trace_id="trace-1",
            model_id="model-a",
            prompt="hello world",
            max_tokens=16,
            estimated_tokens=18,
        )
        prefill = ComputationNode(
            node_id="prefill-1",
            node_type=NodeType.PREFILL,
            created_at=now,
            task_id="task-1",
            operation="prefill",
            input_token_count=2,
            resource_allocation=ResourceAllocation(kv_cache_tokens=2),
            parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL),
        )
        decode = ComputationNode(
            node_id="decode-1",
            node_type=NodeType.DECODE,
            created_at=now,
            task_id="task-1",
            operation="decode",
            input_token_count=1,
            output_token_count=16,
            resource_allocation=ResourceAllocation(kv_cache_tokens=16),
            parallelism=ParallelismConfig(strategy=ParallelismStrategy.TENSOR_PARALLEL),
        )
        ir = SchedulerIR(ir_id="ir-1")
        ir.add_node(task)
        ir.add_node(prefill)
        ir.add_node(decode)
        ir.add_edge(DependencyEdge(from_node="task-1", to_node="prefill-1"))
        ir.add_edge(DependencyEdge(from_node="prefill-1", to_node="decode-1"))
        return ir

    def test_graph_operations(self) -> None:
        """Should support add_node/add_edge/get_node/dependency queries."""
        ir = self._build_sample_ir()
        assert ir.get_node("task-1") is not None
        assert ir.get_node("not-exist") is None
        assert len(ir.get_task_nodes()) == 1
        assert len(ir.get_computation_nodes()) == 2
        assert len(ir.get_dependencies("decode-1")) == 1
        assert len(ir.get_dependents("task-1")) == 1

    def test_json_roundtrip(self) -> None:
        """Should roundtrip IR via JSON serialization."""
        ir = self._build_sample_ir()
        restored = SchedulerIR.from_json(ir.to_json())
        assert restored.ir_id == ir.ir_id
        assert len(restored.nodes) == 3
        assert len(restored.edges) == 2
        assert isinstance(restored.get_node("task-1"), TaskNode)
        assert isinstance(restored.get_node("prefill-1"), ComputationNode)
