"""Unit tests for IRBuilder."""

from __future__ import annotations

import pytest

from sagellm_control.ir.builder import IRBuilder
from sagellm_control.ir.types import NodeType, ParallelismStrategy
from sagellm_control.types import RequestMetadata, RequestPriority
from sagellm_protocol import Request


class TestIRBuilder:
    """Tests for IRBuilder behavior."""

    def test_build_from_protocol_request(self) -> None:
        """Should build IR from Protocol Request."""
        builder = IRBuilder(default_tp_degree=2, default_pp_degree=1)
        request = Request(
            request_id="req-1",
            trace_id="trace-1",
            model="Qwen2-7B",
            prompt="hello world from protocol",
            max_tokens=32,
            stream=False,
        )

        ir = builder.build_from_request(request)

        assert ir.ir_id.startswith("ir-")
        assert len(ir.nodes) == 3
        assert len(ir.edges) == 2
        task = ir.get_node("task-req-1")
        prefill = ir.get_node("prefill-req-1")
        decode = ir.get_node("decode-req-1")
        assert task is not None
        assert prefill is not None
        assert decode is not None

    def test_build_from_request_metadata(self) -> None:
        """Should build IR from RequestMetadata."""
        builder = IRBuilder(enable_kv_reuse=True, kv_reuse_threshold=8)
        metadata = RequestMetadata(
            request_id="req-meta-1",
            trace_id="trace-meta-1",
            model_id="Qwen2-7B",
            prompt="metadata request prompt",
            max_tokens=64,
            prompt_tokens=10,
            priority=RequestPriority.HIGH,
        )

        ir = builder.build_from_request(metadata)

        task = ir.get_node("task-req-meta-1")
        assert task is not None
        assert task.priority == RequestPriority.HIGH.value
        assert task.estimated_tokens == 74
        assert ir.optimization_hints["enable_kv_reuse"] is True
        assert ir.optimization_hints["kv_reuse_threshold"] == 8

    def test_build_from_batch_requests(self) -> None:
        """Should build batch IR from multiple requests."""
        builder = IRBuilder(default_tp_degree=2, default_pp_degree=1)
        requests = [
            RequestMetadata(
                request_id="req-b1",
                trace_id="trace-b1",
                model_id="Qwen2-7B",
                prompt="hello from request one",
                max_tokens=16,
                prompt_tokens=4,
            ),
            RequestMetadata(
                request_id="req-b2",
                trace_id="trace-b2",
                model_id="Qwen2-7B",
                prompt="hello from request two",
                max_tokens=24,
                prompt_tokens=4,
            ),
        ]

        ir = builder.build_from_batch(requests)

        assert ir.ir_id.startswith("ir-batch-")
        assert len(ir.nodes) == 6
        assert len(ir.edges) == 4
        assert ir.global_parallelism.strategy == ParallelismStrategy.DATA_PARALLEL
        assert ir.global_parallelism.dp_degree == 2
        assert ir.optimization_hints["batch_size"] == 2

    def test_build_from_empty_batch_raises(self) -> None:
        """Should reject empty batch."""
        builder = IRBuilder()
        with pytest.raises(ValueError, match="empty batch"):
            builder.build_from_batch([])

    def test_parallelism_strategy_selection(self) -> None:
        """Should select different strategies by prompt length."""
        builder = IRBuilder(default_tp_degree=4, default_pp_degree=1)

        small = builder._select_parallelism_strategy(prompt_tokens=64, max_tokens=64)
        medium = builder._select_parallelism_strategy(prompt_tokens=1024, max_tokens=64)
        large = builder._select_parallelism_strategy(prompt_tokens=4096, max_tokens=64)

        assert small.strategy == ParallelismStrategy.TENSOR_PARALLEL
        assert small.tp_degree == 2
        assert small.pp_degree == 1

        assert medium.strategy == ParallelismStrategy.TENSOR_PARALLEL
        assert medium.tp_degree == 4

        assert large.strategy == ParallelismStrategy.HYBRID
        assert large.pp_degree >= 2

    def test_prefill_and_decode_node_creation(self) -> None:
        """Should create prefill/decode nodes with expected fields."""
        builder = IRBuilder(default_tp_degree=2)
        metadata = RequestMetadata(
            request_id="req-node-1",
            trace_id="trace-node-1",
            model_id="Qwen2-7B",
            prompt="a b c d",
            max_tokens=10,
            prompt_tokens=4,
        )
        ir = builder.build_from_request(metadata)

        prefill = ir.get_node("prefill-req-node-1")
        decode = ir.get_node("decode-req-node-1")
        assert prefill is not None
        assert decode is not None
        assert prefill.node_type == NodeType.PREFILL
        assert decode.node_type == NodeType.DECODE
        assert prefill.resource_allocation.kv_cache_tokens == 4
        assert decode.resource_allocation.kv_cache_tokens == 10

    def test_latency_estimation_functions(self) -> None:
        """Should estimate latency by linear models."""
        builder = IRBuilder()
        assert builder._estimate_prefill_latency(0) == 5.0
        assert builder._estimate_prefill_latency(10) == 10.0
        assert builder._estimate_decode_latency(0) == 10.0
        assert builder._estimate_decode_latency(5) == 60.0
