"""KV cache architecture: three-layer separation (SGLang-inspired).

Layer 1 — KVStorage:      flat token-level K/V buffers [n_layers, size, H, D]
Layer 2 — ReqToTokenPool:  index table [req_idx, pos] → physical token slot
Layer 3 — Allocator:       slot/page allocation with ref-counting and LRU

PagePool orchestrates all three plus RadixCache (prefix addressing).
KVCache is a pure dataclass passed to the model for direct buffer access.

Two modes:
  - contiguous (default): pre-allocated per-request blocks, no dynamic alloc
  - paged: shared pool with on-demand allocation, prefix caching support
"""

import threading
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, OrderedDict

import torch
from torch import Tensor

from astrai.inference.core.workspace import InferenceWorkspace


@dataclass
class _BindState:
    """Cached bind metadata for steady-state decode increment detection."""

    sig: tuple
    seq_lens: List[int]


def _is_steady_increment(
    prev_sig: Optional[tuple],
    prev_vals: Optional[List[int]],
    cur_sig: tuple,
    cur_vals: List[int],
) -> bool:
    """True when the same ordered set has every value +1 from the previous step."""
    return (
        prev_sig is not None
        and prev_vals is not None
        and prev_sig == cur_sig
        and len(prev_vals) == len(cur_vals)
        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
    )


def page_hash(
    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
) -> int:
    start = page_idx * page_size
    end = min(start + page_size, len(token_ids))
    h = parent_hash
    for i in range(start, end):
        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
    return h


class Allocator:
    """Bitmask-based page allocator with ref-counting and LRU eviction."""

    def __init__(self, n_pages: int):
        self._free_mask = (1 << n_pages) - 1
        self._refs: List[int] = [0] * n_pages
        self._lru: OrderedDict[int, None] = OrderedDict()
        self.on_evict: Optional[Callable[[int], None]] = None
        self._lock = threading.Lock()

    def alloc(self) -> int:
        with self._lock:
            if self._free_mask:
                lsb = self._free_mask & -self._free_mask
                idx = lsb.bit_length() - 1
                self._free_mask ^= lsb
                self._refs[idx] = 1
                return idx
            if self._lru:
                idx, _ = self._lru.popitem(last=False)
                if self.on_evict:
                    self.on_evict(idx)
                self._refs[idx] = 1
                self._free_mask &= ~(1 << idx)
                return idx
            return -1

    def free(self, idx: int, keep_cached: bool = False):
        with self._lock:
            self._refs[idx] -= 1
            if self._refs[idx] == 0:
                if keep_cached:
                    self._lru[idx] = None
                else:
                    self._free_mask |= 1 << idx

    def inc_ref(self, idx: int):
        with self._lock:
            self._refs[idx] += 1
            self._lru.pop(idx, None)

    def ref_count(self, idx: int) -> int:
        with self._lock:
            return self._refs[idx]

    def touch(self, idx: int):
        with self._lock:
            if idx in self._lru:
                self._lru.move_to_end(idx)


class RadixNode:
    """A page-aligned edge in the CPU-side prefix radix."""

    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")

    def __init__(self, parent=None, tokens=(), page_idx=None):
        self.parent = parent
        self.children: Dict[tuple, "RadixNode"] = {}
        self.page_idx = page_idx
        self.tokens = tuple(tokens)
        self.lock_ref = 0


class RadixCache:
    """Page-granular radix prefix index with exact token matching."""

    def __init__(self, page_size: int):
        self._page_size = page_size
        self._root = RadixNode()
        self._page_to_node: Dict[int, RadixNode] = {}
        # Retained as an introspection-compatible map; matching never relies on
        # this lossy value.
        self._page_to_hash: Dict[int, int] = {}
        self._lock = threading.Lock()

    def evict(self, idx: int):
        with self._lock:
            node = self._page_to_node.pop(idx, None)
            self._page_to_hash.pop(idx, None)
            if node is None:
                return
            node.page_idx = None
            parent = node.parent
            if parent is not None:
                parent.children.pop(node.tokens, None)

    def has_page(self, idx: int) -> bool:
        with self._lock:
            return idx in self._page_to_node

    def lookup(self, token_ids: List[int]) -> List[int]:
        with self._lock:
            full_pages = len(token_ids) // self._page_size
            hits: List[int] = []
            node = self._root
            for i in range(full_pages):
                start = i * self._page_size
                page_tokens = tuple(token_ids[start : start + self._page_size])
                child = node.children.get(page_tokens)
                if child is None or child.page_idx is None:
                    break
                hits.append(child.page_idx)
                node = child
            return hits

    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
        with self._lock:
            full_pages = len(token_ids) // self._page_size
            if logical_page_idx >= full_pages:
                return
            old = self._page_to_node.pop(page_idx, None)
            self._page_to_hash.pop(page_idx, None)
            if old is not None and old.parent is not None:
                old.parent.children.pop(old.tokens, None)

            node = self._root
            for i in range(logical_page_idx + 1):
                start = i * self._page_size
                page_tokens = tuple(token_ids[start : start + self._page_size])
                child = node.children.get(page_tokens)
                if child is None:
                    child = RadixNode(node, page_tokens)
                    node.children[page_tokens] = child
                node = child
            if node.page_idx is not None and node.page_idx != page_idx:
                replaced = node.page_idx
                self._page_to_node.pop(replaced, None)
                self._page_to_hash.pop(replaced, None)
            node.page_idx = page_idx
            self._page_to_node[page_idx] = node
            self._page_to_hash[page_idx] = page_hash(
                token_ids, logical_page_idx, self._page_size
            )

    def release(self, pages: List[int]) -> None:
        with self._lock:
            for page_idx in pages:
                node = self._page_to_node.get(page_idx)
                if node is not None and node.lock_ref:
                    node.lock_ref -= 1


class ReqToTokenPool:
    """Maps [req_idx, pos] -> physical token slot in KV storage.

    Each row is one request; each column is a sequence position. The value
    at [req_idx, pos] is the flat index into the KV storage buffers.
    """

    def __init__(self, size: int, max_context_len: int, device: torch.device):
        self.size = size
        self.max_context_len = max_context_len
        self.req_to_token = torch.zeros(
            (size, max_context_len), dtype=torch.long, device=device
        )
        self.free_slots = list(range(size))
        self._lock = threading.Lock()

    def alloc(self, num_reqs: int) -> Optional[List[int]]:
        with self._lock:
            if num_reqs > len(self.free_slots):
                return None
            slots = self.free_slots[:num_reqs]
            self.free_slots = self.free_slots[num_reqs:]
            return slots

    def free(self, req_indices: List[int]):
        with self._lock:
            self.free_slots.extend(req_indices)

    def write(self, indices, values):
        self.req_to_token[indices] = values


class KVStorage:
    """Token-level KV cache storage.

    Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies
    one slot indexed by ReqToTokenPool.
    """

    def __init__(
        self,
        size: int,
        n_layers: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.size = size
        self.k_buffer = torch.empty(
            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
        )
        self.v_buffer = torch.empty(
            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
        )

    def get_key_buffer(self, layer_id: int) -> Tensor:
        return self.k_buffer[layer_id]

    def get_value_buffer(self, layer_id: int) -> Tensor:
        return self.v_buffer[layer_id]

    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
        self.k_buffer[layer_id, loc] = k
        self.v_buffer[layer_id, loc] = v


@dataclass
class KVCache:
    """Pure data struct passed to model for KV cache I/O.

    The attention layer does raw buffer indexing — no methods, no abstraction.

    Attributes:
        k_buffer: [n_layers, size, n_kv_heads, head_dim]
        v_buffer: [n_layers, size, n_kv_heads, head_dim]
        req_to_token: [num_reqs, max_ctx_len] — index table
        req_pool_indices: [batch_size] — row indices into req_to_token
        seq_lens: [batch_size] — per-request total sequence lengths
        out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
        max_len: max(seq_lens) as Python int — avoids GPU sync in decode
        kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
            per step so the attention backend avoids rebuilding it per layer.
        qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
        decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
        decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
        decode_out: pre-allocated decode output buffer (graph-safe)
    """

    k_buffer: Tensor
    v_buffer: Tensor
    req_to_token: Tensor
    req_pool_indices: Tensor
    seq_lens: Tensor
    out_cache_loc: Tensor
    max_len: int = 0
    kv_indptr: Optional[Tensor] = None
    qo_indptr: Optional[Tensor] = None
    decode_o_part: Optional[Tensor] = None
    decode_ml_part: Optional[Tensor] = None
    decode_out: Optional[Tensor] = None


class PagePool:
    """Top-level KV cache manager.

    Combines KVStorage + ReqToTokenPool + Allocator + RadixCache.

    Args:
        n_layers: Number of transformer layers.
        n_kv_heads: Number of KV attention heads.
        head_dim: Dimension per head.
        max_batch_size: Maximum concurrent requests.
        max_seq_len: Maximum sequence length per request.
        device, dtype: Tensor device and dtype.
        page_size: Page size for paged mode (1 = token-level).
        n_tokens: Total token slots for paged mode. None = contiguous mode
            (pre-allocates max_batch_size * max_seq_len).
    """

    def __init__(
        self,
        n_layers: int,
        n_kv_heads: int,
        head_dim: int,
        max_batch_size: int,
        max_seq_len: int,
        device: torch.device,
        dtype: torch.dtype,
        page_size: int = 1,
        n_tokens: Optional[int] = None,
    ):
        self.page_size = page_size
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.device = device
        self.dtype = dtype
        self.n_layers = n_layers
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim

        self.contiguous = n_tokens is None
        if self.contiguous:
            self.n_tokens = max_batch_size * max_seq_len
        else:
            self.n_tokens = n_tokens

        self._storage = KVStorage(
            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
        )
        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)

        if self.contiguous:
            for i in range(max_batch_size):
                self._req_pool.req_to_token[i] = torch.arange(
                    i * max_seq_len, (i + 1) * max_seq_len, device=device
                )
            self._alloc: Optional[Allocator] = None
            self._prefix: Optional[RadixCache] = None
        else:
            n_pages = self.n_tokens // page_size
            self._alloc = Allocator(n_pages)
            self._prefix = RadixCache(page_size) if page_size > 1 else None
            if self._prefix is not None:
                self._alloc.on_evict = self._prefix.evict

        self._task_req: Dict[str, int] = {}
        self._task_len: Dict[int, int] = {}
        self._task_cached: Dict[str, int] = {}
        self._task_slots: Dict[str, List[int]] = {}
        self._task_pages: Dict[str, List[int]] = {}
        self._lock = threading.Lock()

        # Steady-state decode validation: when the same ordered task set
        # advances every sequence by exactly one token per step, bind_tasks
        # updates the stable buffers in-place instead of re-cumsumming.
        self._bind_state: Optional[_BindState] = None

    # ---- task lifecycle ----

    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
        req_slots = self._req_pool.alloc(1)
        if req_slots is None:
            return False
        req_idx = req_slots[0]
        self._task_req[task_id] = req_idx

        if self.contiguous:
            self._task_len[req_idx] = len(prompt_ids)
            self._task_cached[task_id] = 0
            return True

        n_tokens_needed = len(prompt_ids)
        cached = 0

        if self._prefix is not None:
            hits = self._prefix.lookup(prompt_ids)
            cached = len(hits) * self.page_size
            for p in hits:
                self._alloc.inc_ref(p)
            self._task_pages[task_id] = list(hits)
            self._task_slots[task_id] = []
        else:
            self._task_pages[task_id] = []
            self._task_slots[task_id] = []

        remaining = n_tokens_needed - cached
        if remaining > 0:
            if self.page_size == 1:
                slots = self._alloc_tokens(remaining)
                if slots is None:
                    for p in self._task_pages[task_id]:
                        self._alloc.free(p)
                    self._req_pool.free([req_idx])
                    del self._task_req[task_id]
                    return False
                self._task_slots[task_id] = slots
            else:
                n_new_pages = (remaining + self.page_size - 1) // self.page_size
                new_pages = []
                for _ in range(n_new_pages):
                    p = self._alloc.alloc()
                    if p < 0:
                        for hp in self._task_pages[task_id]:
                            self._alloc.free(hp)
                        for np_ in new_pages:
                            self._alloc.free(np_)
                        self._req_pool.free([req_idx])
                        del self._task_req[task_id]
                        return False
                    new_pages.append(p)
                self._task_pages[task_id].extend(new_pages)

        self._write_req_to_token(task_id, prompt_ids, cached)
        self._task_len[req_idx] = len(prompt_ids)
        self._task_cached[task_id] = cached
        return True

    def task_free(self, task_id: str):
        req_idx = self._task_req.pop(task_id, None)
        if req_idx is None:
            return
        self._task_len.pop(req_idx, None)
        self._task_cached.pop(task_id, None)

        if not self.contiguous:
            if self._prefix is not None:
                for p in self._task_pages.get(task_id, []):
                    keep = self._prefix.has_page(p)
                    self._alloc.free(p, keep_cached=keep)
                    if not keep:
                        self._prefix.evict(p)
            else:
                for p in self._task_pages.get(task_id, []):
                    self._alloc.free(p)
            self._task_pages.pop(task_id, None)
            self._task_slots.pop(task_id, None)

        self._req_pool.free([req_idx])

    def task_extend(self, task_id: str, pos: int) -> bool:
        req_idx = self._task_req.get(task_id)
        if req_idx is None or pos >= self.max_seq_len:
            return False

        # Paged mode must also claim a physical slot for the new token;
        # contiguous mode's block is pre-allocated so this is a no-op.
        if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
            return False

        self._task_len[req_idx] = pos + 1
        return True

    def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
        """Allocate the physical slot for one extended token (paged mode)."""
        if self.page_size == 1:
            slots = self._alloc_tokens(1)
            if slots is None:
                return False
            self._task_slots.setdefault(task_id, []).extend(slots)
            self._req_pool.req_to_token[req_idx, pos] = slots[0]
            return True

        page_idx = pos // self.page_size
        existing = self._task_pages.get(task_id, [])
        if page_idx >= len(existing):
            p = self._alloc.alloc()
            if p < 0:
                return False
            existing.append(p)
            self._task_pages[task_id] = existing
        page_offset = pos % self.page_size
        page = existing[page_idx]
        token_slot = page * self.page_size + page_offset
        self._req_pool.req_to_token[req_idx, pos] = token_slot
        return True

    def task_cached(self, task_id: str) -> int:
        return self._task_cached.get(task_id, 0)

    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
    ):
        if self._prefix is None or self.contiguous:
            return
        pages = self._task_pages.get(task_id, [])
        full_pages = len(prompt_ids) // self.page_size
        for i in range(start_logical_page, min(full_pages, len(pages))):
            self._prefix.record(pages[i], prompt_ids, i)

    def task_cacheable_ids(
        self, task_id: str, prompt_ids: List[int], output_ids: List[int]
    ):
        """Return the sequence whose KV entries are already materialized.

        The first sampled output is produced by prompt prefill, and the last
        sampled output has not been decoded into KV yet.  Therefore the cache
        can safely retain the prompt plus every output except the last one.
        """
        return list(prompt_ids) + list(output_ids[:-1])

    # ---- bind for forward ----

    def bind_tasks(
        self,
        task_ids: List[str],
        workspace: InferenceWorkspace,
        device: Optional[torch.device] = None,
        start_pos: Optional[int] = None,
    ) -> KVCache:
        if device is None:
            device = workspace.device
        req_indices = [self._task_req[tid] for tid in task_ids]
        # Per-request lengths come from the pool's own tracking (task_alloc
        # sets len(prompt_ids); task_extend sets pos+1), so callers need not
        # pass them.
        seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
        b = len(task_ids)
        sig = tuple(task_ids)

        # Write into the caller's workspace buffers (fixed addresses, sized
        # to max_batch/max_seq at init) — the sole owner of the per-step
        # KV bind tensors.
        rpi_buf = workspace.req_pool_indices
        sl_buf = workspace.seq_lens
        kvp_buf = workspace.kv_indptr
        inc_buf = workspace.inc
        ocl_buf = workspace.out_cache_loc

        prev = self._bind_state
        incremental = (
            start_pos is None
            and prev is not None
            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
        )
        if incremental:
            # Steady-state decode: advance the stable buffers in-place.
            # Normal-mode buffers keep ``+=`` legal regardless of whether
            # this runs inside ``torch.inference_mode()``.
            sl_buf[:b] += 1
            kvp_buf[: b + 1] += inc_buf[: b + 1]
            req_pool_indices = rpi_buf[:b]
            seq_lens_t = sl_buf[:b]
            kv_indptr = kvp_buf[: b + 1]
        else:
            # Cold path: fill the stable buffers from fresh host tensors.
            rpi_buf[:b].copy_(
                torch.tensor(req_indices, dtype=torch.long, device=device)
            )
            sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
            kvp_buf[: b + 1].zero_()
            kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
            req_pool_indices = rpi_buf[:b]
            seq_lens_t = sl_buf[:b]
            kv_indptr = kvp_buf[: b + 1]
        self._bind_state = _BindState(sig, list(seq_lens))

        if start_pos is not None:
            seq_len = seq_lens[0]
            out_cache_loc = self._req_pool.req_to_token[
                req_pool_indices, start_pos:seq_len
            ]
            # Ragged query segmentation for the prefill kernel, computed once
            # (was rebuilt per layer in CudaBackend.fwd_prefill).
            q_len = seq_len - start_pos
            workspace.qo_indptr[: b + 1].copy_(
                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
            )
            qo_indptr = workspace.qo_indptr[: b + 1]
            decode_o_part, decode_ml_part = None, None
            decode_out = None
        else:
            write_pos = seq_lens_t - 1
            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
            ocl_buf[:b].copy_(loc)
            out_cache_loc = ocl_buf[:b]
            qo_indptr = None
            decode_o_part = getattr(workspace, "decode_o_part", None)
            decode_ml_part = getattr(workspace, "decode_ml_part", None)
            decode_out = getattr(workspace, "decode_out", None)

        return KVCache(
            k_buffer=self._storage.k_buffer,
            v_buffer=self._storage.v_buffer,
            req_to_token=self._req_pool.req_to_token,
            req_pool_indices=req_pool_indices,
            seq_lens=seq_lens_t,
            out_cache_loc=out_cache_loc,
            max_len=max(seq_lens),
            kv_indptr=kv_indptr,
            qo_indptr=qo_indptr,
            decode_o_part=decode_o_part,
            decode_ml_part=decode_ml_part,
            decode_out=decode_out,
        )

    # ---- internals ----

    def _alloc_tokens(self, n: int) -> Optional[List[int]]:
        if self.page_size != 1:
            raise RuntimeError("_alloc_tokens is for page_size=1 only")
        slots = []
        for _ in range(n):
            p = self._alloc.alloc()
            if p < 0:
                for s in slots:
                    self._alloc.free(s)
                return None
            slots.append(p)
        return slots

    def _write_req_to_token(self, task_id: str, prompt_ids: List[int], cached: int):
        req_idx = self._task_req[task_id]
        total = len(prompt_ids)

        if self.contiguous:
            return

        if self.page_size == 1:
            slots = self._task_slots.get(task_id, [])
            all_slots = slots[: total - cached]
            if all_slots:
                self._req_pool.req_to_token[req_idx, cached:total] = torch.tensor(
                    all_slots, dtype=torch.long, device=self.device
                )
        else:
            pages = self._task_pages.get(task_id, [])
            for pos in range(cached, total):
                page_idx = pos // self.page_size
                page_offset = pos % self.page_size
                if page_idx < len(pages):
                    token_slot = pages[page_idx] * self.page_size + page_offset
                    self._req_pool.req_to_token[req_idx, pos] = token_slot