"""A paged KV cache for one attention head: block pool, free list, block tables, copy-on-write."""
from collections import deque

import numpy as np


class OutOfBlocks(Exception):
    """A block is needed and the free list is empty."""


class BlockManager:
    def __init__(self, num_blocks: int, block_size: int, head_dim: int = 1):
        self.block_size = block_size
        # The pool: all physical memory, cut into fixed-size blocks, allocated once.
        self.k_pool = np.zeros((num_blocks, block_size, head_dim))
        self.v_pool = np.zeros((num_blocks, block_size, head_dim))
        self.free_list = deque(range(num_blocks))   # take from the left, give back on the right
        self.ref_count = [0] * num_blocks           # how many block tables point at each block
        self.tables = {}                            # request id -> physical block ids, in logical order
        self.lengths = {}                           # request id -> tokens stored

    def add_request(self, rid) -> None:
        self.tables[rid] = []
        self.lengths[rid] = 0                       # allocates nothing

    def _take_block(self) -> int:
        block = self.free_list.popleft()
        self.ref_count[block] = 1
        return block

    def append(self, rid, k: np.ndarray, v: np.ndarray) -> None:
        """Store the key and value of one new token at the end of request `rid`."""
        table, n = self.tables[rid], self.lengths[rid]
        offset = n % self.block_size
        last_is_full = offset == 0                  # also true when the table is empty
        last_is_shared = not last_is_full and self.ref_count[table[-1]] > 1
        if (last_is_full or last_is_shared) and not self.free_list:
            raise OutOfBlocks(rid)                  # nothing has changed yet
        if last_is_full:
            table.append(self._take_block())
        elif last_is_shared:                        # copy-on-write: give this request its own copy
            old, new = table[-1], self._take_block()
            self.k_pool[new], self.v_pool[new] = self.k_pool[old], self.v_pool[old]
            self.ref_count[old] -= 1
            table[-1] = new
        self.k_pool[table[-1], offset] = k
        self.v_pool[table[-1], offset] = v
        self.lengths[rid] = n + 1

    def fork(self, parent, child) -> None:
        """The child starts as the parent's twin. No data is copied: both tables name the same blocks."""
        self.tables[child] = list(self.tables[parent])
        self.lengths[child] = self.lengths[parent]
        for block in self.tables[child]:
            self.ref_count[block] += 1

    def free(self, rid) -> None:
        """Finish a request. A block returns to the free list when no table points at it any more."""
        for block in self.tables.pop(rid):
            self.ref_count[block] -= 1
            if self.ref_count[block] == 0:
                self.free_list.append(block)
        del self.lengths[rid]

    def locate(self, rid, position: int):
        """Logical token position -> (physical block, slot inside the block)."""
        return self.tables[rid][position // self.block_size], position % self.block_size

    def gather(self, rid):
        """All keys and values of `rid` in logical order, shape [tokens, head_dim]."""
        n, head_dim = self.lengths[rid], self.k_pool.shape[2]
        keys, values = np.zeros((n, head_dim)), np.zeros((n, head_dim))
        for pos in range(n):
            block, slot = self.locate(rid, pos)
            keys[pos], values[pos] = self.k_pool[block, slot], self.v_pool[block, slot]
        return keys, values


def paged_attention(q: np.ndarray, mgr: BlockManager, rid) -> np.ndarray:
    """One query vector attends over request `rid`, reading the cache one block at a time."""
    n, size, table = mgr.lengths[rid], mgr.block_size, mgr.tables[rid]
    scores = []
    for j, block in enumerate(table):
        valid = min(size, n - j * size)             # the last block may be partly filled
        scores.append(mgr.k_pool[block, :valid] @ q)
    scores = np.concatenate(scores) / np.sqrt(len(q))
    weights = np.exp(scores - scores.max())
    weights /= weights.sum()
    out = np.zeros_like(q, dtype=float)
    for j, block in enumerate(table):
        valid = min(size, n - j * size)
        out += weights[j * size : j * size + valid] @ mgr.v_pool[block, :valid]
    return out


def attention_contiguous(q: np.ndarray, keys: np.ndarray, values: np.ndarray) -> np.ndarray:
    """Reference: the same attention over keys and values already laid out in one array."""
    scores = keys @ q / np.sqrt(len(q))
    weights = np.exp(scores - scores.max())
    return (weights / weights.sum()) @ values


def worked_example():
    """The lesson's hand-worked example: block size 4, a pool of 8 blocks, requests A, B and C.

    Returns a list of (event, tables, free_list, lengths), one per step.
    """
    mgr = BlockManager(num_blocks=8, block_size=4)
    steps = []

    def grow(rid, tokens):
        for _ in range(tokens):
            n = mgr.lengths[rid]
            mgr.append(rid, np.array([n]), np.array([n]))

    def snap(event):
        steps.append((event, {r: list(t) for r, t in mgr.tables.items()}, list(mgr.free_list),
                      dict(mgr.lengths)))

    snap("start")
    for rid, prompt in (("A", 6), ("B", 3), ("C", 5)):
        mgr.add_request(rid)
        grow(rid, prompt)
        snap(f"{rid} arrives with a {prompt}-token prompt")
    grow("A", 2)
    snap("A generates 2 tokens")
    grow("A", 1)
    snap("A generates 1 more token")
    grow("B", 2)
    snap("B generates 2 tokens")
    mgr.free("B")
    snap("B finishes")
    grow("C", 4)
    snap("C generates 4 tokens")
    return steps


def prefix_sharing_example():
    """Block size 4, a 6-token prompt shared by two samples. Returns (event, tables, ref_counts, free_list)."""
    mgr = BlockManager(num_blocks=6, block_size=4)
    steps = []

    def snap(event):
        steps.append((event, {r: list(t) for r, t in mgr.tables.items()}, list(mgr.ref_count),
                      list(mgr.free_list)))

    def one(rid):
        n = mgr.lengths[rid]
        mgr.append(rid, np.array([n]), np.array([n]))

    mgr.add_request("S1")
    for _ in range(6):
        one("S1")
    snap("S1 holds the 6-token prompt")
    mgr.fork("S1", "S2")
    snap("S2 forks from S1")
    one("S1")
    snap("S1 writes token 6")
    one("S2")
    snap("S2 writes token 6")
    return steps