"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache.

Layout
    Two pinned host pools, both organized in pages:
      - KV pages:    one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block.
      - State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a
                     128-token-aligned position with GPUCacheInterface.select_state_slice_to_store.
    Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)),
    so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool.

Transfer
    "kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through
    UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never
    blocks the CPU. Completion is tracked with CUDA events polled from the engine loop.

Policy
    Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last
    128-aligned position are stored (pages already present are only touched). On a new request, the longest
    page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to
    be recomputed, then snapped down to the nearest position that has a stored state; the sequence is
    pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with
    num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is
    rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store).

Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's
prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank).
"""

from __future__ import annotations

import hashlib
import os
import time
from collections import OrderedDict, deque
from dataclasses import dataclass, field

import numpy as np
import torch

from deepinfer import envs
from deepinfer.logger import init_logger
from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface

logger = init_logger(__name__)


class _Entry:
    __slots__ = ("key", "slot", "refs", "ready")

    def __init__(self, key: bytes, slot: int) -> None:
        self.key = key
        self.slot = slot
        self.refs = 0  # in-flight transfers (store or load) touching this slot
        self.ready = False  # host copy complete and hit-able


@dataclass
class _Chain:
    tokens: np.ndarray
    hashes: list[bytes] = field(default_factory=list)
    counted: bool = False  # lookup stats recorded (a deferred request is looked up repeatedly)


@dataclass
class Stats:
    lookups: int = 0
    hits: int = 0
    hit_tokens: int = 0
    lookup_tokens: int = 0
    deferred: int = 0  # hit found but GPU allocation not possible yet
    stores: int = 0
    store_skipped: int = 0
    pages_stored: int = 0
    pages_loaded: int = 0
    states_stored: int = 0
    states_loaded: int = 0
    evicted_pages: int = 0
    evicted_states: int = 0
    errors: int = 0

    def as_dict(self) -> dict[str, int]:
        return dict(self.__dict__)


class HostPrefixCache:
    def __init__(self, engine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
        assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)"
        assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1"
        assert GPUCacheInterface.cache_info_ready, "cache engine not initialized"

        self.engine = engine
        self.scheduler = engine.scheduler
        cache_engine = engine.driver_worker.cache_engine
        self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
        self.gpu_states: torch.Tensor = cache_engine.gpu_states
        assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8
        self.device = self.gpu_blocks.device
        self.block_size: int = engine.cache_config.block_size
        assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN)
        self.page_bytes: int = self.gpu_blocks.size(1)
        self.state_bytes: int = GPUCacheInterface.state_size_to_disk
        assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes)
        self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS
        # The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative.
        self.min_hit_tokens = 2 * self.block_size
        # A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only
        # at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a
        # block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA
        # ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore.
        self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4"))
        self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference

        self.num_pages = max(0, size_bytes // self.page_bytes)
        self.num_state_slots = num_state_slots
        self.staging_slots = staging_slots
        assert self.num_pages > 0 and num_state_slots > 0
        self.gpu_blocks_ptr = self.gpu_blocks.data_ptr()
        # Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be
        # cached before a prefill completes, and an unused cache then costs no pinned memory).
        self._allocated = False
        self.host_pages = self.host_states = None
        self.host_pages_ptr = self.host_states_ptr = 0
        self._staging: list[torch.Tensor] = []
        self._staging_idx = 0
        self.kernel = None

        self.pages: OrderedDict[bytes, _Entry] = OrderedDict()
        self.states: OrderedDict[bytes, _Entry] = OrderedDict()
        self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1))
        self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1))
        self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque()
        self._chains: dict[str, _Chain] = {}
        self.stats = Stats()
        self.enabled = True
        # DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel).
        self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1"
        logger.info(
            "HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d",
            self.num_pages,
            self.page_bytes,
            self.num_pages * self.page_bytes / 2**30,
            num_state_slots,
            self.state_bytes,
            num_state_slots * self.state_bytes / 2**30,
            self.keep_tokens,
        )

    def _allocate(self) -> None:
        if self._allocated:
            return
        from deepinfer.hicache.kernels import copy_pages_module

        t0 = time.perf_counter()
        self.kernel = copy_pages_module()
        self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True)
        self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True)
        self.host_pages_ptr = self.host_pages.data_ptr()
        self.host_states_ptr = self.host_states.data_ptr()
        # Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream.
        self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)]
        self._allocated = True
        logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0)

    # ------------------------------------------------------------------ hashing

    def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None:
        chain = self._chains.get(request_id)
        if chain is None and tokens is not None:
            chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32))
        return chain

    def _extend(self, chain: _Chain, num_pages: int) -> None:
        hashes, toks, bs = chain.hashes, chain.tokens, self.block_size
        prev = hashes[-1] if hashes else b""
        for i in range(len(hashes), num_pages):
            prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest()
            hashes.append(prev)

    def forget(self, request_id: str) -> None:
        self._chains.pop(request_id, None)

    # ------------------------------------------------------------------ transfers

    def _ids(self, ids: list[int]) -> torch.Tensor:
        return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True)

    def _next_staging(self) -> torch.Tensor:
        t = self._staging[self._staging_idx]
        self._staging_idx = (self._staging_idx + 1) % len(self._staging)
        return t

    def _track(self, entries: list[_Entry]) -> None:
        ev = torch.cuda.Event()
        ev.record()
        self._inflight.append((ev, entries))

    def _checkpoint(self, what: str) -> None:
        if self._debug_sync:
            from deepinfer.utils import allow_cuda_sync

            with allow_cuda_sync():
                torch.cuda.synchronize()
            logger.info("hicache debug: ok after %s", what)

    def poll(self) -> None:
        """Retire completed transfers (events on one stream complete in order, so only the head is checked)."""
        while self._inflight and self._inflight[0][0].query():
            _, entries = self._inflight.popleft()
            for e in entries:
                e.refs -= 1
                e.ready = True

    def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int:
        """Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found."""
        victims, busy = [], []
        for key, e in pool.items():
            if len(victims) >= need:
                break
            if e.refs or not e.ready:
                busy.append(key)
                continue
            victims.append(key)
        for key in busy:
            pool.move_to_end(key)
        if len(victims) < need:
            return 0
        for key in victims:
            free.append(pool.pop(key).slot)
        return len(victims)

    def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool:
        if need > len(free) + len(pool):
            return False
        if len(free) < need:
            evicted = self._evict(pool, free, need - len(free))
            if pool is self.pages:
                self.stats.evicted_pages += evicted
            else:
                self.stats.evicted_states += evicted
        return len(free) >= need

    # ------------------------------------------------------------------ load (hit)

    def try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
        """Fill `seq` (not yet in the scheduler) from the host cache.

        Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not
        allocatable right now, sequence untouched, retry later).
        """
        if not self.enabled:
            return 0
        try:
            return self._try_resume(seq, request_id, prompt_ids)
        except Exception:  # noqa: BLE001 - never take the engine down for a cache problem
            logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id)
            self._fail()
            return 0

    def _try_resume(self, seq, request_id: str, prompt_ids: np.ndarray) -> int:
        bs = self.block_size
        n = len(prompt_ids)
        chain = self._chain(request_id, prompt_ids)
        if not chain.counted:
            chain.counted = True
            self.stats.lookups += 1
            self.stats.lookup_tokens += n
        limit = ((n - self.keep_tokens) // bs) * bs  # the replay rectangle [n - keep, n) must be recomputed
        if limit < self.min_hit_tokens:
            return 0
        max_pages = limit // bs
        self._extend(chain, max_pages)
        matched = 0
        for h in chain.hashes[:max_pages]:
            e = self.pages.get(h)
            if e is None or not e.ready:
                break
            matched += 1
        b = min(matched * bs, limit)
        while b >= self.min_hit_tokens:
            st = self.states.get(chain.hashes[b // bs - 1])
            if st is not None and st.ready:
                break
            b -= bs
        else:
            return 0

        try:
            allocated = self.scheduler.try_allocate_seq(seq)
        except RuntimeError:  # can never fit (prompt longer than the whole KV pool): let the scheduler reject it
            return 0
        if not allocated:
            self.stats.deferred += 1
            return -1
        num_pages = b // bs
        page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)]
        for e in page_entries:
            e.refs += 1
            self.pages.move_to_end(e.key)
        st.refs += 1
        self.states.move_to_end(st.key)

        block_table = seq.block_table
        self.kernel.copy_pages(
            self.host_pages_ptr,
            self.gpu_blocks_ptr,
            self._ids([e.slot for e in page_entries]),
            self._ids(list(block_table[:num_pages])),
            self.page_bytes,
        )
        self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})")
        staging = self._next_staging()
        self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes)
        self._checkpoint("load: state copy H2D")
        GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states)
        self._checkpoint("load: load_state_to_cache")
        self._track(page_entries + [st])

        seq.num_cached_tokens = b
        seq.cached_hit_tokens = b
        self.stats.hits += 1
        self.stats.hit_tokens += b
        self.stats.pages_loaded += num_pages
        self.stats.states_loaded += 1
        return b

    # ------------------------------------------------------------------ store (write-through at prefill end)

    def on_prefill_step(self, scheduled, finished) -> None:
        """Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed."""
        if not self.enabled:
            return
        try:
            for seq in scheduled:
                if seq.chunked_rest_len == 0:  # last chunk of this prompt
                    self._store(seq)
        except Exception:  # noqa: BLE001
            logger.exception("host prefix cache: store failed")
            self._fail()

    def _store(self, seq) -> None:
        chain = self._chains.get(seq.request_id)
        if chain is None:  # request did not go through try_resume: hash the prompt from the sequence itself
            chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32))
        self._allocate()
        bs = self.block_size
        b = (seq.get_len() // bs) * bs
        if b < self.min_hit_tokens:
            return
        num_pages = b // bs
        self._extend(chain, num_pages)
        missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages]
        for i in range(num_pages):
            if i not in missing:
                self.pages.move_to_end(chain.hashes[i])
        # State boundaries: b, b - bs, ... while the window ending there is still in the ring.
        seq_len = seq.get_len()
        new_states = []
        for k in range(self.state_lookback_blocks):
            pos = b - k * bs
            if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack:
                break
            key = chain.hashes[pos // bs - 1]
            if key in self.states:
                self.states.move_to_end(key)
            else:
                new_states.append((pos, key))
        if not missing and not new_states:
            return
        if not self._reserve(self.pages, self.free_pages, len(missing)):
            self.stats.store_skipped += 1
            return
        # States are best effort: keep the newest boundaries that fit.
        while new_states and not self._reserve(self.states, self.free_states, len(new_states)):
            new_states.pop()
        if not missing and not new_states:
            self.stats.store_skipped += 1
            return

        block_table = seq.block_table
        entries: list[_Entry] = []
        if missing:
            new = []
            for i in missing:
                e = _Entry(chain.hashes[i], self.free_pages.pop())
                e.refs = 1
                self.pages[e.key] = e
                new.append(e)
            self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})")
            self.kernel.copy_pages(
                self.gpu_blocks_ptr,
                self.host_pages_ptr,
                self._ids([block_table[i] for i in missing]),
                self._ids([e.slot for e in new]),
                self.page_bytes,
            )
            self._checkpoint("store: page copy D2H")
            entries += new
            self.stats.pages_stored += len(new)
        for pos, key in new_states:
            st = _Entry(key, self.free_states.pop())
            st.refs = 1
            self.states[key] = st
            staging = self._next_staging()
            GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging)
            self._checkpoint(f"store: select_state_slice at {pos}")
            self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes)
            self._checkpoint("store: state copy D2H")
            entries.append(st)
            self.stats.states_stored += 1
        self._track(entries)
        self.stats.stores += 1

    # ------------------------------------------------------------------ misc

    def _fail(self) -> None:
        self.stats.errors += 1
        if self.stats.errors >= 3:
            logger.error("host prefix cache: disabled after %d errors", self.stats.errors)
            self.enabled = False

    def stats_dict(self) -> dict[str, int | float]:
        d = self.stats.as_dict()
        d.update(
            enabled=int(self.enabled),
            pages_used=len(self.pages),
            pages_total=self.num_pages,
            states_used=len(self.states),
            states_total=len(self.free_states) + len(self.states),
            inflight=len(self._inflight),
            hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0,
        )
        return d