from collections import deque
import xxhash
import numpy as np

from nanovllm.engine.sequence import Sequence


class Block:
    """KV cache 的物理块。

    每个 Block 对应 GPU 显存中一块固定大小的 KV cache 存储，可以存放 block_size 个 token 的 KV 向量。

    Attributes:
        block_id: 物理块的唯一 ID。
        ref_count: 引用计数，有多少个序列正在共享这个块（前缀缓存复用）。
        hash: 该块内容的哈希值，用于前缀缓存查找。初始为 -1 表示未计算哈希。
        token_ids: 该块对应的 token ID 列表，用于验证缓存命中时内容是否一致。
    """

    def __init__(self, block_id):
        self.block_id = block_id
        self.ref_count = 0
        self.hash = -1
        self.token_ids = []

    def update(self, hash: int, token_ids: list[int]):
        """更新块的内容哈希和 token_ids。"""
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        """重置块状态，准备被重新分配。"""
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


class BlockManager:
    """管理 KV cache 物理块的分配、释放和前缀缓存。

    核心设计:
    - 将 KV cache 分成固定大小的物理块（block_size 个 token/块）。
    - 使用哈希表实现前缀缓存（prefix caching）：不同序列如果拥有相同的 prompt 前缀，
      可以共享同一组 KV cache 块，避免重复计算。
    - 块通过引用计数（ref_count）管理生命周期，所有引用释放后块变为空闲可复用。

    Attributes:
        block_size: 每个块存储的 token 数量。
        blocks: 所有物理块的列表，索引即 block_id。
        hash_to_block_id: 内容哈希 → 块 ID 的映射，用于前缀缓存查找。
        free_block_ids: 空闲块的 ID 队列。
        used_block_ids: 正在使用的块 ID 集合。
    """

    def __init__(self, num_blocks: int, block_size: int):
        self.block_size = block_size
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
        self.hash_to_block_id: dict[int, int] = dict()
        self.free_block_ids: deque[int] = deque(range(num_blocks))
        self.used_block_ids: set[int] = set()

    @classmethod
    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
        """计算一个块的内容哈希值。

        哈希是链式的：每个块的哈希值依赖于前一个块的哈希值（prefix 参数），
        这确保了前缀缓存的一致性——只有完全相同的 token 序列前缀才会产生相同的哈希链。

        Args:
            token_ids: 该块对应的 token ID 列表。
            prefix: 前一个块的哈希值，-1 表示第一个块（无前缀）。
        """
        h = xxhash.xxh64()
        if prefix != -1:
            h.update(prefix.to_bytes(8, "little"))
        h.update(np.array(token_ids).tobytes())
        return h.intdigest()

    def _allocate_block(self) -> int:
        """从空闲池中分配一个物理块。"""
        block_id = self.free_block_ids.popleft()
        block = self.blocks[block_id]
        assert block.ref_count == 0
        # 如果该块之前有哈希记录，先从哈希表中移除
        if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
            del self.hash_to_block_id[block.hash]
        block.reset()
        self.used_block_ids.add(block_id)
        return block_id

    def _deallocate_block(self, block_id: int):
        """释放一个物理块回空闲池。"""
        assert self.blocks[block_id].ref_count == 0
        self.used_block_ids.remove(block_id)
        self.free_block_ids.append(block_id)

    def can_allocate(self, seq: Sequence) -> int:
        """检查是否有足够的空闲块来分配给序列，同时计算前缀缓存命中数。

        遍历序列的所有逻辑块，逐个计算哈希并与已有缓存比对：
        - 如果哈希匹配且 token 内容一致，则该块可以复用（前缀缓存命中）。
        - 一旦遇到不匹配的块就停止，因为前缀缓存要求连续匹配。
        - 最后检查剩余空闲块是否足够分配未命中部分。

        Returns:
            前缀缓存命中的块数，如果没有足够的空闲块则返回 -1。
        """
        h = -1
        num_cached_blocks = 0
        num_new_blocks = seq.num_blocks
        for i in range(seq.num_blocks - 1):
            token_ids = seq.block(i)
            h = self.compute_hash(token_ids, h)
            block_id = self.hash_to_block_id.get(h, -1)
            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
                break
            num_cached_blocks += 1
            if block_id in self.used_block_ids:
                # 块已在使用中（被其他序列共享），不需要额外分配
                num_new_blocks -= 1
        if len(self.free_block_ids) < num_new_blocks:
            return -1
        return num_cached_blocks

    def allocate(self, seq: Sequence, num_cached_blocks: int):
        """为序列分配 KV cache 物理块。

        前缀缓存命中部分：增加已有块的引用计数。
        未命中部分：从空闲池分配新块。

        Args:
            seq: 要分配块的序列。
            num_cached_blocks: 前缀缓存命中的块数（由 can_allocate 计算）。
        """
        assert not seq.block_table
        h = -1
        # 处理缓存命中部分：复用已有块
        for i in range(num_cached_blocks):
            token_ids = seq.block(i)
            h = self.compute_hash(token_ids, h)
            block_id = self.hash_to_block_id[h]
            block = self.blocks[block_id]
            if block_id in self.used_block_ids:
                block.ref_count += 1
            else:
                # 块在空闲池中但哈希匹配（之前被释放但内容未覆盖），重新激活
                block.ref_count = 1
                self.free_block_ids.remove(block_id)
                self.used_block_ids.add(block_id)
            seq.block_table.append(block_id)
        # 分配新块用于未命中部分
        for i in range(num_cached_blocks, seq.num_blocks):
            seq.block_table.append(self._allocate_block())
        seq.num_cached_tokens = num_cached_blocks * self.block_size

    def deallocate(self, seq: Sequence):
        """释放序列占用的所有 KV cache 块。

        递减每个块的引用计数，引用计数归零的块被回收进空闲池。
        注意：遍历顺序为逆序，这是为了优先释放最新分配的块（它们在空闲队列尾部）。
        """
        for block_id in reversed(seq.block_table):
            block = self.blocks[block_id]
            block.ref_count -= 1
            if block.ref_count == 0:
                self._deallocate_block(block_id)
        seq.num_cached_tokens = 0
        seq.block_table.clear()

    def can_append(self, seq: Sequence) -> bool:
        """检查 decode 阶段是否能追加一个 token。

        当序列当前最后一个块已满时（len(seq) % block_size == 1，即新增 token 是新块的第一个），
        需要分配一个新块。否则只需要写入已有块，不需要额外分配。
        """
        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

    def may_append(self, seq: Sequence):
        """在 decode 阶段，如果需要则追加一个新块。"""
        if len(seq) % self.block_size == 1:
            seq.block_table.append(self._allocate_block())

    def hash_blocks(self, seq: Sequence):
        """更新序列中已完成计算的块的哈希值。

        在每次调度步骤完成后调用，将新计算完毕的块的 token 内容和哈希值
        注册到哈希表中，以便后续序列复用（前缀缓存）。

        只处理从 num_cached_tokens 到当前进度的块（即本次新完成的块）。
        """
        start = seq.num_cached_tokens // self.block_size
        end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size
        if start == end: return
        # 继承前一个块的哈希值作为前缀
        h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1
        for i in range(start, end):
            block = self.blocks[seq.block_table[i]]
            token_ids = seq.block(i)
            h = self.compute_hash(token_ids, h)
            block.update(h, token_ids)
            self.hash_to_block_id[h] = block.block_id