"""TTS 预训练权重 safetensors 迁移。

transformers >= 4.56 的安全策略 (CVE-2025-32434) 拒绝在 torch < 2.6 上加载
pytorch_model.bin。transformers 加载时优先使用 model.safetensors，因此首次
启动时把 .bin 转换一次即可，无需升级 torch。
"""
from pathlib import Path

from voice_dialogue.config import paths
from voice_dialogue.utils.logger import logger

PRETRAINED_DIRS = [
    "chinese-roberta-wwm-ext-large",
    "chinese-hubert-base",
]


def ensure_safetensors_weights() -> None:
    """确保 MoYoYo TTS 的预训练权重存在 safetensors 版本，缺失时从 .bin 转换。"""
    moyoyo_path = Path(paths.TTS_MODELS_PATH) / "moyoyo"

    for dirname in PRETRAINED_DIRS:
        model_dir = moyoyo_path / dirname
        bin_path = model_dir / "pytorch_model.bin"
        st_path = model_dir / "model.safetensors"

        if st_path.exists() or not bin_path.exists():
            continue

        logger.info(f"[INFO] 首次启动：转换 {dirname} 权重为 safetensors...")
        try:
            import torch
            from safetensors.torch import save_file

            state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)
            # clone 断开共享内存，safetensors 不允许张量间共享存储
            state_dict = {
                key: value.clone().contiguous()
                for key, value in state_dict.items()
                if hasattr(value, "clone")
            }
            save_file(state_dict, st_path, metadata={"format": "pt"})
            logger.info(f"[INFO] {dirname} 转换完成: {st_path.stat().st_size // 1024 ** 2} MB")
        except Exception as e:
            logger.error(f"[ERROR] 转换 {dirname} 权重失败: {e}")