"""Tests for KV hooks extension (issue #18 / #19)"""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from sagellm_protocol import (
    KVAllocateParams,
    KVDType,
    KVEvictParams,
    KVEvictReason,
    KVFreeParams,
    KVHandle,
    KVHookEvent,
    KVHookPayload,
    KVLayout,
    KVMigrateParams,
    KVPrefetchParams,
    KVPriority,
    KVQoSPolicy,
)


class TestKVPriority:
    def test_all_priorities_exist(self) -> None:
        assert KVPriority.CRITICAL == "critical"
        assert KVPriority.HIGH == "high"
        assert KVPriority.MEDIUM == "medium"
        assert KVPriority.LOW == "low"
        assert KVPriority.BACKGROUND == "background"


class TestKVQoSPolicy:
    def test_defaults(self) -> None:
        qos = KVQoSPolicy()
        assert qos.tier == "default"
        assert qos.slo_latency_ms is None
        assert qos.max_wait_ms is None
        assert qos.evict_protected is False

    def test_premium_tier(self) -> None:
        qos = KVQoSPolicy(
            tier="premium",
            slo_latency_ms=50.0,
            max_wait_ms=200.0,
            evict_protected=True,
        )
        assert qos.tier == "premium"
        assert qos.evict_protected is True

    def test_slo_latency_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            KVQoSPolicy(slo_latency_ms=0.0)

    def test_max_wait_ms_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            KVQoSPolicy(max_wait_ms=-1.0)


class TestKVAllocateParamsExtended:
    """Tests for extended KVAllocateParams fields (issue #18/#19)."""

    def test_backward_compatible_original_fields(self) -> None:
        """Original callers without new fields still work."""
        params = KVAllocateParams(
            request_id="req-001",
            num_tokens=512,
            kv_dtype=KVDType.FP16,
            layout=KVLayout.BSHD,
            block_size=64,
        )
        assert params.request_id == "req-001"
        assert params.priority == 0
        assert params.trace_id is None
        assert params.qos_labels == {}

    def test_with_priority_and_trace(self) -> None:
        params = KVAllocateParams(
            request_id="req-002",
            num_tokens=1024,
            kv_dtype=KVDType.BF16,
            layout=KVLayout.BHSD,
            block_size=128,
            priority=10,
            trace_id="trace-abc",
            qos_labels={"tier": "premium"},
        )
        assert params.priority == 10
        assert params.trace_id == "trace-abc"
        assert params.qos_labels["tier"] == "premium"

    def test_metadata_field(self) -> None:
        params = KVAllocateParams(
            request_id="req-003",
            num_tokens=256,
            kv_dtype=KVDType.INT8,
            layout=KVLayout.BSHD,
            block_size=32,
            metadata={"engine_id": "engine-0"},
        )
        assert params.metadata["engine_id"] == "engine-0"


class TestKVHookEvent:
    def test_all_events_exist(self) -> None:
        events = [
            KVHookEvent.ALLOCATE_BEFORE,
            KVHookEvent.ALLOCATE_AFTER,
            KVHookEvent.FREE_BEFORE,
            KVHookEvent.FREE_AFTER,
            KVHookEvent.EVICT_BEFORE,
            KVHookEvent.EVICT_AFTER,
            KVHookEvent.MIGRATE_BEFORE,
            KVHookEvent.MIGRATE_AFTER,
            KVHookEvent.PREFETCH_BEFORE,
            KVHookEvent.PREFETCH_AFTER,
        ]
        assert len(events) == 10


class TestKVHookPayload:
    def _make_handle(self) -> KVHandle:
        return KVHandle(
            handle_id="handle-001",
            request_id="req-001",
            num_tokens=512,
            device="cuda:0",
        )

    def test_allocate_after_payload(self) -> None:
        payload = KVHookPayload(
            event=KVHookEvent.ALLOCATE_AFTER,
            request_id="req-001",
            trace_id="trace-001",
            kv_handle=self._make_handle(),
            device="cuda:0",
            elapsed_ms=0.5,
            success=True,
        )
        assert payload.event == KVHookEvent.ALLOCATE_AFTER
        assert payload.success is True
        assert payload.kv_handle is not None

    def test_evict_payload_with_reason(self) -> None:
        payload = KVHookPayload(
            event=KVHookEvent.EVICT_BEFORE,
            request_id="req-002",
            kv_handle=self._make_handle(),
            evict_reason=KVEvictReason.MEMORY_PRESSURE,
        )
        assert payload.evict_reason == KVEvictReason.MEMORY_PRESSURE

    def test_failed_payload(self) -> None:
        payload = KVHookPayload(
            event=KVHookEvent.ALLOCATE_AFTER,
            request_id="req-003",
            success=False,
            error_message="out of memory",
        )
        assert payload.success is False
        assert "out of memory" in (payload.error_message or "")

    def test_default_values(self) -> None:
        payload = KVHookPayload(
            event=KVHookEvent.FREE_AFTER,
            request_id="req-004",
        )
        assert payload.kv_handle is None
        assert payload.evict_reason is None
        assert payload.elapsed_ms == 0.0
        assert payload.success is True
        assert payload.metadata == {}


class TestKVFreeParams:
    def _make_handle(self) -> KVHandle:
        return KVHandle(
            handle_id="handle-001",
            request_id="req-001",
            num_tokens=512,
            device="cuda:0",
        )

    def test_basic_free(self) -> None:
        params = KVFreeParams(kv_handle=self._make_handle())
        assert params.force is False
        assert params.trace_id is None

    def test_force_free_with_trace(self) -> None:
        params = KVFreeParams(
            kv_handle=self._make_handle(),
            force=True,
            trace_id="trace-001",
        )
        assert params.force is True


class TestKVEvictParams:
    def _make_handle(self) -> KVHandle:
        return KVHandle(
            handle_id="handle-001",
            request_id="req-001",
            num_tokens=512,
            device="cuda:0",
        )

    def test_evict_to_cpu(self) -> None:
        params = KVEvictParams(
            kv_handle=self._make_handle(),
            reason=KVEvictReason.MEMORY_PRESSURE,
            target_device="cpu",
        )
        assert params.reason == KVEvictReason.MEMORY_PRESSURE
        assert params.target_device == "cpu"

    def test_evict_discard(self) -> None:
        params = KVEvictParams(
            kv_handle=self._make_handle(),
            reason=KVEvictReason.EXPLICIT_FREE,
        )
        assert params.target_device is None

    def test_ttl_evict(self) -> None:
        params = KVEvictParams(
            kv_handle=self._make_handle(),
            reason=KVEvictReason.TTL_EXPIRED,
            priority=5,
        )
        assert params.reason == KVEvictReason.TTL_EXPIRED
        assert params.priority == 5


class TestKVMigrateParamsExtended:
    """Tests for extended KVMigrateParams (new optional fields)."""

    def _make_handle(self) -> KVHandle:
        return KVHandle(
            handle_id="handle-001",
            request_id="req-001",
            num_tokens=512,
            device="cuda:0",
        )

    def test_migrate_with_trace_and_priority(self) -> None:
        params = KVMigrateParams(
            kv_handle=self._make_handle(),
            target_device="cuda:1",
            trace_id="trace-001",
            priority=3,
        )
        assert params.target_device == "cuda:1"
        assert params.trace_id == "trace-001"
        assert params.priority == 3

    def test_migrate_backward_compat(self) -> None:
        """Original callers without trace/priority still work."""
        params = KVMigrateParams(
            kv_handle=self._make_handle(),
            target_device="npu:0",
        )
        assert params.trace_id is None
        assert params.priority == 0


class TestKVPrefetchParams:
    def _make_handle(self) -> KVHandle:
        return KVHandle(
            handle_id="handle-001",
            request_id="req-001",
            num_tokens=512,
            device="cpu",
        )

    def test_basic_prefetch(self) -> None:
        params = KVPrefetchParams(
            request_id="req-001",
            kv_handle=self._make_handle(),
            target_device="cuda:0",
        )
        assert params.target_device == "cuda:0"
        assert params.priority == 0
        assert params.trace_id is None
        assert params.prefetch_window_tokens is None

    def test_prefetch_with_window(self) -> None:
        params = KVPrefetchParams(
            request_id="req-001",
            kv_handle=self._make_handle(),
            target_device="cuda:0",
            prefetch_window_tokens=256,
            priority=5,
        )
        assert params.prefetch_window_tokens == 256
        assert params.priority == 5

    def test_prefetch_window_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            KVPrefetchParams(
                request_id="req-001",
                kv_handle=self._make_handle(),
                target_device="cuda:0",
                prefetch_window_tokens=0,
            )
