"""ComfyUI 节点: MiniMax H3 Adapter Loader（假 CLIP 注入）。

用法（T2V 工作流）:
    UNETLoader(fl2va) -> model
    MiniMaxH3AdapterLoader(student=<Qwen3.5-4B 文件夹>, adapter=<adapter.safetensors>) -> clip
    VAELoader(video_vae) -> vae
    MiniMaxH3ImageToVideo(clip=clip, vae=vae, prompt=..., width, height, length) -> positive + latent
    BasicGuider(model, positive) -> RandomNoise -> SamplerCustomAdvanced -> VAEDecode -> SaveVideo

原理:
    ComfyUI 的 MiniMaxH3 DiT (comfy/ldm/minimax/model.py) 在
    preprocess_text_embeds / _forward 中条件式跳过投影:
        if text_states.shape[-1] != self.hidden_size:  # 5376
            text_states = token_refiner(condition_proj(text_states))
    适配器输出已是 [1, S_T, 5376]（post-refiner），DiT 直接使用，零手术。

    官方 MiniMaxH3ImageToVideo 只对 clip 调用 tokenize(prompt) +
    encode_from_tokens_scheduled(tokens)，故用鸭子类型 CLIP 对象替换即可，
    latent 创建 / keyframe / duration 网格全部保留。
"""
from __future__ import annotations

import os

import torch

import folder_paths
import comfy.utils
import comfy.model_management as mm

from .adapter.model import H3Adapter
from .adapter_clip import MiniMaxH3AdapterCLIP, MiniMaxH3AdapterFromCLIP
from .student import StudentTextEncoder, get_torch_device

ADAPTER_FOLDER = "minimax_h3_adapter"   # models/minimax_h3_adapter/
STUDENT_FOLDER = "minimax_h3_student"   # -> models/text_encoders/


def _h3_tokenizer() -> object:
    from transformers import AutoTokenizer
    tok_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tokenizer")
    return AutoTokenizer.from_pretrained(tok_dir)


def _student_folder_options() -> list[str]:
    """models/text_encoders/ 下的子目录（HF 格式）或 .gguf 文件（transformers GGUF 加载）。"""
    opts = []
    for root in folder_paths.get_folder_paths(STUDENT_FOLDER):
        if os.path.isdir(root):
            opts += [d for d in sorted(os.listdir(root))
                     if os.path.isdir(os.path.join(root, d)) and os.path.exists(os.path.join(root, d, "config.json"))]
            opts += [f for f in sorted(os.listdir(root)) if f.endswith(".gguf")]
    return opts or ["<把 Qwen3.5-4B 文件夹或 .gguf 放到 models/text_encoders/>"]


class MiniMaxH3AdapterLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "student": (_student_folder_options(),),
            "adapter": (folder_paths.get_filename_list(ADAPTER_FOLDER),),
        }, "optional": {
            # 留空 = 整卡；24GB 及以下填 "16GiB" 等启用层间 offload
            "gpu_mem": ("STRING", {"default": ""}),
        }}

    RETURN_TYPES = ("CLIP",)
    RETURN_NAMES = ("clip",)
    FUNCTION = "load_adapter"
    CATEGORY = "model/conditioning/minimax"

    def load_adapter(self, student: str, adapter: str, gpu_mem: str = ""):
        dev = get_torch_device()
        model = _load_adapter_weights(adapter, dev)

        student_root = folder_paths.get_folder_paths(STUDENT_FOLDER)[0]
        student_dir = os.path.join(student_root, student)
        stu = StudentTextEncoder(student_dir, lowvram=True, gpu_mem=gpu_mem)

        tok = _h3_tokenizer()
        clip = MiniMaxH3AdapterCLIP(model, stu, tok, dev)
        return (clip,)


def _load_adapter_weights(adapter: str, dev: torch.device) -> H3Adapter:
    """加载适配器 safetensors -> bf16 H3Adapter（容错 model. 前缀）。"""
    adapter_path = folder_paths.get_full_path(ADAPTER_FOLDER, adapter)
    sd = comfy.utils.load_torch_file(adapter_path)
    if any(k.startswith("model.") for k in sd):
        sd = {k[len("model."):]: v for k, v in sd.items()}
    model = H3Adapter().to(torch.bfloat16)
    missing, unexpected = model.load_state_dict(sd, strict=False)
    assert not missing, f"adapter 缺 key: {missing[:10]}"
    if unexpected:
        print(f"    [adapter] ignored {len(unexpected)} unexpected keys", flush=True)
    model.to(dev).eval()
    for m in model.modules():
        if hasattr(m, "use_sdpa"):
            m.use_sdpa = True
    print(f"    [adapter] loaded {adapter} ({sum(p.numel() for p in model.parameters())/1e9:.3f}B)", flush=True)
    return model


class MiniMaxH3AdapterFromCLIPLoader:
    """CLIP→CLIP 包装: 学生 CLIP（CLIPLoader GGUF 等）的 encode 输出经适配器 -> 可喂给官方 MiniMaxH3ImageToVideo。

    工作流:
        CLIPLoader(GGUF, 学生模型) -> clip
        MiniMaxH3AdapterFromCLIPLoader(clip=clip, adapter=adapter_stage2.safetensors) -> clip'
        MiniMaxH3ImageToVideo(clip=clip', ...)  # 官方节点原样
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "clip": ("CLIP",),
            "adapter": (folder_paths.get_filename_list(ADAPTER_FOLDER),),
        }}

    RETURN_TYPES = ("CLIP",)
    RETURN_NAMES = ("clip",)
    FUNCTION = "wrap_clip"
    CATEGORY = "model/conditioning/minimax"

    def wrap_clip(self, clip, adapter):
        dev = get_torch_device()
        model = _load_adapter_weights(adapter, dev)
        wrapped = MiniMaxH3AdapterFromCLIP(clip, model, _h3_tokenizer(), dev)
        return (wrapped,)