"""
Paged Video Latent 内存池与连续批处理引擎 (vLLM Style)
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""

import torch
import logging
from typing import Dict, List

logger = logging.getLogger("VerseFlow.PagedLatentPool")

class LatentBlock:
    def __init__(self, block_id: int, block_size: int, num_channels: int, height: int, width: int, dtype=torch.float16):
        self.block_id = block_id
        self.ref_count = 0
        self.data = torch.empty((block_size, num_channels, height, width), dtype=dtype, device="cuda")

class PagedVideoLatentMemoryPool:
    """视频扩散模型 Paged Latent 显存池，消除显存碎片"""
    def __init__(self, total_blocks: int = 1024, block_size: int = 16, channels: int = 4, height: int = 64, width: int = 64):
        self.block_size = block_size
        self.free_blocks: List[LatentBlock] = [
            LatentBlock(i, block_size, channels, height, width) for i in range(total_blocks)
        ]
        self.allocated_blocks: Dict[int, LatentBlock] = {}
        self.block_table: Dict[str, List[int]] = {}

    def allocate(self, sequence_id: str, num_frames: int) -> List[int]:
        num_blocks_needed = (num_frames + self.block_size - 1) // self.block_size
        if len(self.free_blocks) < num_blocks_needed:
            raise RuntimeError(f"Paged 显存池开销溢出: 需要 {num_blocks_needed} 页，剩余 {len(self.free_blocks)} 页")

        allocated_ids = []
        for _ in range(num_blocks_needed):
            block = self.free_blocks.pop()
            block.ref_count = 1
            self.allocated_blocks[block.block_id] = block
            allocated_ids.append(block.block_id)

        self.block_table[sequence_id] = allocated_ids
        logger.info(f"🧠 [PagedLatent] 序列 {sequence_id} 成功分配 {num_blocks_needed} 个 Block 页")
        return allocated_ids

    def free_sequence(self, sequence_id: str):
        if sequence_id not in self.block_table:
            return

        for block_id in self.block_table[sequence_id]:
            block = self.allocated_blocks.pop(block_id, None)
            if block:
                block.ref_count = 0
                self.free_blocks.append(block)

        del self.block_table[sequence_id]
        logger.info(f"♻️ [PagedLatent] 序列 {sequence_id} 显存页成功回收")