"""Integration-oriented tests for KVCacheManager."""

from __future__ import annotations

import pytest

from sagellm_kv_cache.errors import KVBudgetExceededError
from sagellm_kv_cache.kv_cache_manager import KVCacheManager
from sagellm_kv_cache.observability import MetricsCollector


class TestKVCacheManagerPrefix:
    def test_register_and_query_prefix_hit(self) -> None:
        manager = KVCacheManager(max_tokens=1024, block_size=16)
        handle = manager.alloc(num_tokens=64)

        manager.register_prefix((10, 11, 12, 13), handle)
        matched_handle, matched_tokens = manager.query_prefix((10, 11, 12, 13, 14))

        assert matched_handle is not None
        assert matched_handle.handle_id == handle.handle_id
        assert matched_tokens == 4
        assert handle.is_pinned

    def test_query_prefix_reused_blocks(self) -> None:
        manager = KVCacheManager(max_tokens=1024, block_size=16)
        handle = manager.alloc(num_tokens=32)
        manager.register_prefix((1, 2, 3, 4), handle)

        blocks, matched_tokens = manager.get_reused_blocks((1, 2, 3, 4, 5))

        assert matched_tokens == 4
        assert len(blocks) > 0

    def test_prefix_metrics_hit_and_miss(self) -> None:
        collector = MetricsCollector(trace_id="trace-1", request_id="req-1", engine_id="engine-0")
        manager = KVCacheManager(max_tokens=1024, block_size=16, metrics_collector=collector)
        handle = manager.alloc(num_tokens=64)
        manager.register_prefix((1, 2, 3), handle)

        manager.query_prefix((1, 2, 3, 4))
        manager.query_prefix((9, 9, 9))

        metrics = manager.get_prefix_metrics()
        assert metrics["prefix_hits"] == 1
        assert metrics["prefix_lookups"] == 2
        assert metrics["prefix_hit_rate"] == 0.5
        assert metrics["prefix_reuse_tokens"] == 3

        assert collector.get_gauge("prefix_hit_rate") == 0.5
        assert collector.get_gauge("prefix_reuse_tokens") == 3.0

    def test_free_invalidates_prefix_mapping(self) -> None:
        manager = KVCacheManager(max_tokens=1024, block_size=16)
        handle = manager.alloc(num_tokens=64)
        manager.register_prefix((1, 2, 3), handle)

        manager.free(handle)
        matched_handle, matched_tokens = manager.query_prefix((1, 2, 3, 4))

        assert matched_handle is None
        assert matched_tokens == 0


class TestKVCacheManagerEviction:
    def test_alloc_auto_evicts_by_lru(self) -> None:
        manager = KVCacheManager(max_tokens=300)

        h1 = manager.alloc(100, metadata={"request_id": "req-1"})
        h2 = manager.alloc(100, metadata={"request_id": "req-2"})
        h3 = manager.alloc(100, metadata={"request_id": "req-3"})

        h1.last_access = 1.0
        h2.last_access = 2.0
        h3.last_access = 3.0

        h4 = manager.alloc(150, metadata={"request_id": "req-4"})

        remaining_ids = set(manager.kv_pool.handles.keys())
        assert h1.handle_id not in remaining_ids
        assert h2.handle_id not in remaining_ids
        assert h3.handle_id in remaining_ids
        assert h4.handle_id in remaining_ids

    def test_pinned_handle_is_not_evicted(self) -> None:
        manager = KVCacheManager(max_tokens=300)

        h1 = manager.alloc(100, metadata={"request_id": "req-prefix"})
        manager.register_prefix([1, 2, 3], h1)

        h2 = manager.alloc(100, metadata={"request_id": "req-2"})
        h3 = manager.alloc(100, metadata={"request_id": "req-3"})

        h1.last_access = 1.0
        h2.last_access = 2.0
        h3.last_access = 3.0

        manager.alloc(150, metadata={"request_id": "req-4"})

        remaining_ids = set(manager.kv_pool.handles.keys())
        assert h1.handle_id in remaining_ids

    def test_all_pinned_raises_budget_error(self) -> None:
        manager = KVCacheManager(max_tokens=100)

        h1 = manager.alloc(100, metadata={"request_id": "req-1"})
        manager.register_prefix([9, 9, 9], h1)

        with pytest.raises(KVBudgetExceededError):
            manager.alloc(10, metadata={"request_id": "req-2"})

    def test_records_evict_metrics(self) -> None:
        collector = MetricsCollector(
            trace_id="trace-kv-1",
            request_id="req-kv-1",
            engine_id="engine-kv-1",
        )
        manager = KVCacheManager(max_tokens=200, metrics_collector=collector)

        h1 = manager.alloc(100, metadata={"request_id": "req-1"})
        h2 = manager.alloc(100, metadata={"request_id": "req-2"})

        h1.last_access = 1.0
        h2.last_access = 2.0

        manager.alloc(100, metadata={"request_id": "req-3"})

        assert collector.get_counter("evict_count") >= 1
        assert collector.get_gauge("evict_ms") >= 0.0

    def test_respects_single_round_eviction_limit(self) -> None:
        manager = KVCacheManager(max_tokens=300, max_evict_handles_per_round=1)

        h1 = manager.alloc(100)
        h2 = manager.alloc(100)
        h3 = manager.alloc(100)

        h1.last_access = 1.0
        h2.last_access = 2.0
        h3.last_access = 3.0

        with pytest.raises(KVBudgetExceededError):
            manager.alloc(150)
