# SPDX-License-Identifier: Apache-2.0
"""Tests for the LMCache cache simulator tool."""

# Standard
from pathlib import Path
import json
import tempfile

# Third Party
import pytest

# First Party
from lmcache.tools.cache_simulator.lru_cache import LRUCache, LRUCacheFast
from lmcache.tools.cache_simulator.simulator import (
    compute_kv_bytes_per_chunk,
    load_lookup_events,
    simulate,
)

# ---------------------------------------------------------------------------
# LRUCacheFast
# ---------------------------------------------------------------------------


class TestLRUCacheFast:
    def test_basic_insert_and_contains(self):
        cache = LRUCacheFast(3)
        cache.insert("a")
        assert cache.contains("a")
        assert not cache.contains("b")

    def test_evicts_lru_when_full(self):
        cache = LRUCacheFast(2)
        cache.insert("a")
        cache.insert("b")
        cache.insert("c")  # should evict "a"
        assert not cache.contains("a")
        assert cache.contains("b")
        assert cache.contains("c")
        assert cache.eviction_count == 1

    def test_access_refreshes_lru_order(self):
        cache = LRUCacheFast(2)
        cache.insert("a")
        cache.insert("b")
        cache.access("a")  # "a" is now MRU; "b" is LRU
        cache.insert("c")  # should evict "b"
        assert cache.contains("a")
        assert not cache.contains("b")
        assert cache.contains("c")

    def test_insert_existing_key_does_not_grow(self):
        cache = LRUCacheFast(2)
        cache.insert("a")
        cache.insert("a")
        assert len(cache) == 1
        assert cache.eviction_count == 0

    def test_capacity_one(self):
        cache = LRUCacheFast(1)
        cache.insert("a")
        cache.insert("b")
        assert not cache.contains("a")
        assert cache.contains("b")
        assert cache.eviction_count == 1

    def test_invalid_capacity_raises(self):
        with pytest.raises(ValueError):
            LRUCacheFast(0)


# ---------------------------------------------------------------------------
# LRUCache
# ---------------------------------------------------------------------------


class TestLRUCache:
    def test_basic_insert_and_contains(self):
        cache = LRUCache(3)
        cache.insert("a")
        assert cache.contains("a")
        assert not cache.contains("z")

    def test_evicts_lru_when_full(self):
        cache = LRUCache(2)
        cache.insert("a")
        cache.insert("b")
        cache.insert("c")  # evicts "a"
        assert not cache.contains("a")
        assert cache.eviction_count == 1

    def test_position_mru_is_zero(self):
        cache = LRUCache(3)
        cache.insert("a")
        cache.insert("b")
        cache.insert("c")
        # "c" was inserted last → MRU → position 0
        assert cache.position("c") == 0
        # "a" was inserted first → LRU → position 2
        assert cache.position("a") == 2

    def test_access_moves_to_mru(self):
        cache = LRUCache(3)
        cache.insert("a")
        cache.insert("b")
        cache.insert("c")
        cache.access("a")  # "a" becomes MRU
        assert cache.position("a") == 0
        assert cache.position("c") == 1
        assert cache.position("b") == 2

    def test_invalid_capacity_raises(self):
        with pytest.raises(ValueError):
            LRUCache(0)


# ---------------------------------------------------------------------------
# compute_kv_bytes_per_chunk
# ---------------------------------------------------------------------------


class TestComputeKvBytesPerChunk:
    def test_float16(self):
        event = {"shapes": [[32, 256, 128]], "dtypes": ["float16"]}
        assert compute_kv_bytes_per_chunk(event) == 32 * 256 * 128 * 2

    def test_multiple_tensors(self):
        # Two tensors: one float16, one float8_e4m3fn
        event = {
            "shapes": [[4, 256, 64], [4, 256, 64]],
            "dtypes": ["float16", "float8_e4m3fn"],
        }
        expected = 4 * 256 * 64 * 2 + 4 * 256 * 64 * 1
        assert compute_kv_bytes_per_chunk(event) == expected

    def test_empty_shapes_returns_zero(self):
        assert compute_kv_bytes_per_chunk({"shapes": [], "dtypes": []}) == 0

    def test_missing_fields_returns_zero(self):
        assert compute_kv_bytes_per_chunk({}) == 0


# ---------------------------------------------------------------------------
# load_lookup_events
# ---------------------------------------------------------------------------


class TestLoadLookupEvents:
    def _write_jsonl(self, lines):
        f = tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False, dir="/tmp")
        for line in lines:
            f.write(json.dumps(line) + "\n")
        f.close()
        return Path(f.name)

    def test_loads_and_sorts_by_timestamp(self):
        path = self._write_jsonl(
            [
                {
                    "timestamp": 2.0,
                    "chunk_hashes": ["0xbb"],
                    "seq_len": 256,
                    "chunk_size": 256,
                },
                {
                    "timestamp": 1.0,
                    "chunk_hashes": ["0xaa"],
                    "seq_len": 256,
                    "chunk_size": 256,
                },
            ]
        )
        events = load_lookup_events([path])
        assert events[0]["chunk_hashes"] == ["0xaa"]
        assert events[1]["chunk_hashes"] == ["0xbb"]
        path.unlink()

    def test_model_filter(self):
        path = self._write_jsonl(
            [
                {
                    "timestamp": 1.0,
                    "model_name": "A",
                    "chunk_hashes": [],
                    "seq_len": 0,
                    "chunk_size": 256,
                },
                {
                    "timestamp": 2.0,
                    "model_name": "B",
                    "chunk_hashes": [],
                    "seq_len": 0,
                    "chunk_size": 256,
                },
            ]
        )
        events = load_lookup_events([path], model="A")
        assert len(events) == 1
        assert events[0]["model_name"] == "A"
        path.unlink()

    def test_max_samples(self):
        path = self._write_jsonl(
            [
                {
                    "timestamp": float(i),
                    "chunk_hashes": [],
                    "seq_len": 0,
                    "chunk_size": 256,
                }
                for i in range(10)
            ]
        )
        events = load_lookup_events([path], max_samples=3)
        assert len(events) == 3
        path.unlink()

    def test_skips_malformed_lines(self):
        f = tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False, dir="/tmp")
        f.write(
            '{"timestamp": 1.0, "chunk_hashes": [], "seq_len": 0, "chunk_size": 256}\n'
        )
        f.write("not json\n")
        f.write(
            '{"timestamp": 2.0, "chunk_hashes": [], "seq_len": 0, "chunk_size": 256}\n'
        )
        f.close()
        # Standard
        import warnings

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            events = load_lookup_events([Path(f.name)])
        assert len(events) == 2
        Path(f.name).unlink()


# ---------------------------------------------------------------------------
# simulate
# ---------------------------------------------------------------------------


class TestSimulate:
    """Core simulation correctness tests."""

    def _make_events(self, specs):
        """Build minimal event dicts from (chunk_hashes, seq_len) pairs."""
        return [
            {"chunk_hashes": hashes, "seq_len": seq_len, "chunk_size": 256}
            for hashes, seq_len in specs
        ]

    def test_cold_cache_zero_hits(self):
        events = self._make_events(
            [
                (["0xaa", "0xbb"], 600),
            ]
        )
        res = simulate(events, cache_capacity_bytes=100, kv_bytes_per_chunk=1)
        assert res["token_hit_rate"] == 0.0
        assert res["total_hit_tokens"] == 0
        assert res["total_tokens"] == 600

    def test_second_identical_request_hits(self):
        # r1: cold miss; r2: full prefix hit (2 chunks × 256 = 512 tokens)
        events = self._make_events(
            [
                (["0xaa", "0xbb"], 600),
                (["0xaa", "0xbb"], 600),
            ]
        )
        res = simulate(events, cache_capacity_bytes=100, kv_bytes_per_chunk=1)
        assert res["total_hit_tokens"] == 512
        assert res["total_tokens"] == 1200
        assert abs(res["token_hit_rate"] - 512 / 1200) < 1e-9

    def test_prefix_semantics_breaks_on_first_miss(self):
        # r1 populates 0xaa, 0xbb, 0xcc
        # r2: 0xaa hits, 0xdd misses → prefix=1; 0xbb is in cache but not counted
        events = self._make_events(
            [
                (["0xaa", "0xbb", "0xcc"], 768),
                (["0xaa", "0xdd", "0xbb"], 768),
            ]
        )
        res = simulate(events, cache_capacity_bytes=100, kv_bytes_per_chunk=1)
        # Only 0xaa is a prefix hit in r2 → 1 × 256 = 256 hit tokens
        assert res["total_hit_tokens"] == 256

    def test_tail_tokens_always_miss(self):
        # seq_len=300, chunk_size=256 → 1 full chunk + 44 tail tokens
        # On second request the full chunk hits but tail is always a miss
        events = self._make_events(
            [
                (["0xaa"], 300),
                (["0xaa"], 300),
            ]
        )
        res = simulate(events, cache_capacity_bytes=100, kv_bytes_per_chunk=1)
        # hit tokens = 256, total = 600 → tail (44 tokens per request) is miss
        assert res["total_hit_tokens"] == 256
        assert res["total_tokens"] == 600

    def test_eviction_reduces_hits(self):
        # capacity = 1 chunk; insert 0xaa then 0xbb evicts 0xaa
        events = self._make_events(
            [
                (["0xaa"], 256),  # populates 0xaa
                (["0xbb"], 256),  # evicts 0xaa, populates 0xbb
                (["0xaa"], 256),  # 0xaa was evicted → miss; evicts 0xbb
            ]
        )
        res = simulate(events, cache_capacity_bytes=1, kv_bytes_per_chunk=1)
        assert res["total_hit_tokens"] == 0
        assert res["eviction_count"] == 2

    def test_fast_mode_matches_normal_token_hit_rate(self):
        events = self._make_events(
            [
                (["0xaa", "0xbb"], 600),
                (["0xaa", "0xbb"], 600),
                (["0xcc"], 300),
            ]
        )
        cap, bpc = 100, 1
        res_normal = simulate(events, cap, bpc, fast=False)
        res_fast = simulate(events, cap, bpc, fast=True)
        assert abs(res_normal["token_hit_rate"] - res_fast["token_hit_rate"]) < 1e-9

    def test_invalid_kv_bytes_raises(self):
        with pytest.raises(ValueError):
            simulate([], cache_capacity_bytes=100, kv_bytes_per_chunk=0)