"""HostPrefixCache logic against a fake engine: hashing/matching, the keep-tokens cap, snapping to a stored
state, store dedupe, LRU eviction and deferral. The real state slice kernels are replaced by plain copies;
the page transfer kernel is the real one. Needs one GPU."""

from __future__ import annotations

import os

import numpy as np
import pytest
import torch

os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = "140"
os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"

pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")

PAGE = 114048
STATE = 4096
BLOCK = 128


class FakeSeq:
    def __init__(self, request_id: str, tokens: np.ndarray) -> None:
        self.request_id = request_id
        self.tokens = tokens
        self.block_table: list[int] = []
        self.state_id = -1
        self.num_cached_tokens = 0
        self.cached_hit_tokens = 0
        self.chunked_rest_len = 0
        self.data = self  # SequenceData stand-in

    def get_len(self) -> int:
        return len(self.tokens)

    def get_token_ids(self) -> list[int]:
        return self.tokens.tolist()


class FakeScheduler:
    def __init__(self, num_blocks: int, num_states: int) -> None:
        self.free_blocks = list(range(num_blocks))
        self.free_states = list(range(num_states))
        self.refuse = False

    def try_allocate_seq(self, seq: FakeSeq) -> bool:
        n = (len(seq.tokens) + BLOCK - 1) // BLOCK
        if self.refuse or len(self.free_blocks) < n or not self.free_states:
            return False
        seq.block_table = [self.free_blocks.pop(0) for _ in range(n)]
        seq.state_id = self.free_states.pop(0)
        return True

    def get_num_unfinished_seqs(self) -> int:
        return 0


class _Obj:
    pass


@pytest.fixture
def engine(monkeypatch):
    from deepinfer.worker.cache_engine import GPUCacheInterface

    monkeypatch.setattr(GPUCacheInterface, "cache_info_ready", True, raising=False)
    monkeypatch.setattr(GPUCacheInterface, "state_size_to_disk", STATE, raising=False)
    monkeypatch.setattr(GPUCacheInterface, "swa_block_size_kvstore", 134 + 512 + 140, raising=False)
    monkeypatch.setattr(GPUCacheInterface, "swa_block_size_inference", 134, raising=False)
    monkeypatch.setattr(
        GPUCacheInterface,
        "select_state_slice_to_store",
        classmethod(lambda cls, sid, pos, cache, state_tensor=None, TO_CE=False: state_tensor.copy_(cache[sid, :STATE])),
    )
    monkeypatch.setattr(
        GPUCacheInterface,
        "load_state_from_store_to_prefill_cache",
        classmethod(lambda cls, sid, pos, tensor, cache: cache[sid, :STATE].copy_(tensor)),
    )
    e = _Obj()
    e.scheduler = FakeScheduler(num_blocks=512, num_states=8)
    e.cache_config = _Obj()
    e.cache_config.block_size = BLOCK
    e.driver_worker = _Obj()
    e.driver_worker.cache_engine = _Obj()
    e.driver_worker.cache_engine.gpu_blocks = torch.zeros((512, PAGE), dtype=torch.uint8, device="cuda")
    e.driver_worker.cache_engine.gpu_states = torch.zeros((8, STATE), dtype=torch.uint8, device="cuda")
    return e


def _prefill(engine, cache, seq: FakeSeq) -> None:
    """Simulate the engine: allocate (if the cache did not), write random KV/state, run the store hook."""
    if not seq.block_table:
        assert engine.scheduler.try_allocate_seq(seq)
    blocks = engine.driver_worker.cache_engine.gpu_blocks
    states = engine.driver_worker.cache_engine.gpu_states
    start = seq.num_cached_tokens // BLOCK
    for b in seq.block_table[start:]:
        blocks[b] = torch.randint(0, 256, (PAGE,), dtype=torch.uint8, device="cuda")
    states[seq.state_id] = torch.randint(0, 256, (STATE,), dtype=torch.uint8, device="cuda")
    cache.on_prefill_step([seq], [])
    torch.cuda.synchronize()
    cache.poll()


def _tokens(n: int, seed: int) -> np.ndarray:
    return np.random.default_rng(seed).integers(0, 100000, n, dtype=np.int32)


def test_hit_snaps_to_stored_state_and_restores_bytes(engine):
    from deepinfer.hicache import HostPrefixCache

    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8)
    blocks = engine.driver_worker.cache_engine.gpu_blocks
    states = engine.driver_worker.cache_engine.gpu_states

    a = FakeSeq("a", _tokens(1000, 1))
    assert cache.try_resume(a, "a", a.tokens) == 0  # cold
    _prefill(engine, cache, a)
    # 896 tokens = 7 pages; states at the last 4 boundaries (896, 768, 640, 512)
    assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4
    a_pages = [blocks[b].clone() for b in a.block_table[:7]]
    a_state = states[a.state_id, :STATE].clone()

    # Same prefix plus a new turn: hit at 896 (= 7 full pages, and >= keep_tokens remain to recompute).
    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
    assert cache.try_resume(b, "b", b.tokens) == 896
    torch.cuda.synchronize()
    cache.poll()
    assert b.num_cached_tokens == 896 and b.cached_hit_tokens == 896
    assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(7))
    assert torch.equal(states[b.state_id, :STATE], a_state)
    assert cache.stats.hits == 1 and cache.stats.hit_tokens == 896

    # Same 950-token prefix but a different tail at 1000 tokens: pages match up to 896, the keep cap allows
    # only 768, and a state exists there -> hit at 768 (the lookback states make short extensions hit).
    c = FakeSeq("c", np.concatenate([a.tokens[:950], _tokens(50, 3)]))
    assert cache.try_resume(c, "c", c.tokens) == 768
    torch.cuda.synchronize()
    cache.poll()
    _prefill(engine, cache, c)  # pages 0..6 and all four states already exist: nothing new is stored
    assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4

    # b's prefill stores its remaining pages (1500 -> 11 pages, 4 new) and states at 1408..1024.
    _prefill(engine, cache, b)
    assert cache.stats.pages_stored == 11 and cache.stats.states_stored == 8
    d = FakeSeq("d", np.concatenate([b.tokens, _tokens(300, 4)]))
    assert cache.try_resume(d, "d", d.tokens) == 1408


def test_eviction_and_capacity(engine):
    from deepinfer.hicache import HostPrefixCache

    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2)
    a = FakeSeq("a", _tokens(7 * BLOCK + 10, 1))
    _prefill(engine, cache, a)
    assert len(cache.pages) == 7

    big = FakeSeq("big", _tokens(60 * BLOCK, 2))  # does not fit at all
    _prefill(engine, cache, big)
    assert cache.stats.store_skipped == 1 and len(cache.pages) == 7

    e = FakeSeq("e", _tokens(40 * BLOCK, 3))  # fits without eviction (43 free)
    _prefill(engine, cache, e)
    assert cache.stats.evicted_pages == 0 and len(cache.pages) == 47

    f = FakeSeq("f", _tokens(20 * BLOCK, 4))  # needs 20, 3 free -> evict 17 LRU pages (a's 7 first)
    _prefill(engine, cache, f)
    assert cache.stats.evicted_pages == 17 and len(cache.pages) == 50
    assert (
        cache.try_resume(FakeSeq("a2", np.concatenate([a.tokens, _tokens(300, 5)])), "a2", np.concatenate([a.tokens, _tokens(300, 5)])) == 0
    )
    # Only two state slots: each store keeps the two newest boundaries that fit, evicting LRU (a's, then e's).
    assert len(cache.states) == 2 and cache.stats.evicted_states == 4


def test_deferred_when_gpu_cannot_allocate(engine):
    from deepinfer.hicache import HostPrefixCache

    cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2)
    a = FakeSeq("a", _tokens(1000, 1))
    _prefill(engine, cache, a)
    engine.scheduler.refuse = True
    b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)]))
    assert cache.try_resume(b, "b", b.tokens) == -1
    assert b.block_table == [] and b.num_cached_tokens == 0 and cache.stats.deferred == 1
    engine.scheduler.refuse = False
    assert cache.try_resume(b, "b", b.tokens) == 896