import os
from dataclasses import dataclass
from transformers import AutoConfig


@dataclass(slots=True)
class Config:
    """推理引擎的全局配置。

    Args:
        model: HuggingFace 模型的本地路径（必须是一个目录，包含 safetensors 权重文件和 config.json）。
        max_num_batched_tokens: 单次调度（schedule）中允许的最大 token 总数，控制 prefill 阶段的批处理粒度。
        max_num_seqs: 单次调度中允许的最大序列数，限制同时处理的请求数量。
        max_model_len: 模型支持的最大序列长度（prompt + 生成的总长度），会被 HF config 中的 max_position_embeddings 截断。
        gpu_memory_utilization: GPU 显存使用率（0~1），决定 KV cache 可用显存大小。
        tensor_parallel_size: 张量并行度（TP），即使用多少张 GPU 来并行运行模型。
        enforce_eager: 是否强制使用 eager 模式（不使用 CUDA Graph），调试时设为 True。
        hf_config: 从模型目录自动加载的 HuggingFace 配置对象，由 __post_init__ 自动设置。
        eos: End-of-Sequence token 的 ID，由引擎初始化时从 tokenizer 获取，用于判断生成是否结束。
        kvcache_block_size: KV cache 的块大小（token 数），每个 block 存储这么多 token 的 KV 向量。前缀缓存和 KV cache 管理的最小单位。
        num_kvcache_blocks: KV cache 的总块数，-1 表示尚未分配，会在 ModelRunner.allocate_kv_cache() 中根据可用显存自动计算。
    """

    model: str
    max_num_batched_tokens: int = 16384
    max_num_seqs: int = 512
    max_model_len: int = 4096
    gpu_memory_utilization: float = 0.9
    tensor_parallel_size: int = 1
    enforce_eager: bool = False
    hf_config: AutoConfig | None = None
    eos: int = -1
    kvcache_block_size: int = 256
    num_kvcache_blocks: int = -1

    def __post_init__(self):
        assert os.path.isdir(self.model)
        assert self.kvcache_block_size % 256 == 0
        assert 1 <= self.tensor_parallel_size <= 8
        # 从模型目录加载 HuggingFace 配置（如 num_layers, hidden_size, num_heads 等）
        self.hf_config = AutoConfig.from_pretrained(self.model)
        # 确保最大序列长度不超过模型支持的位置编码上限
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)