# nodes.py

import logging
import os

import comfy.lora
import comfy.model_base
import comfy.utils
import folder_paths

from .mapping import map_fastvideo_module, split_payload_key


_LORA_A = ".lora_A.weight"
_LORA_B = ".lora_B.weight"
_DIFF = ".diff"
_DIFF_B = ".diff_b"
_SET_WEIGHT = ".set_weight"
_ALPHA = ".alpha"


def _analyze_payload(keys):
    payloads = {}
    for key in keys:
        module, suffix = split_payload_key(key)
        payloads.setdefault(module, set()).add(suffix)

    expected_patches = 0
    for module, suffixes in payloads.items():
        has_a = _LORA_A in suffixes
        has_b = _LORA_B in suffixes
        if has_a != has_b:
            raise ValueError(f"FastH3 adapter has an incomplete LoRA pair for {module}")
        if _ALPHA in suffixes and not has_a:
            raise ValueError(f"FastH3 adapter has alpha without a LoRA pair for {module}")

        weight_payloads = int(has_a) + int(_DIFF in suffixes) + int(_SET_WEIGHT in suffixes)
        if weight_payloads > 1:
            raise ValueError(f"FastH3 adapter has conflicting weight payloads for {module}")
        if weight_payloads == 0 and _DIFF_B not in suffixes:
            raise ValueError(f"FastH3 adapter has no applicable payload for {module}")
        expected_patches += weight_payloads + int(_DIFF_B in suffixes)
    return payloads, expected_patches


def _validate_adapter(metadata, payloads):
    adapter_format = (metadata or {}).get("format")
    if adapter_format is not None and adapter_format != "fastvideo-lora-v2":
        raise ValueError(
            f"Expected a FastVideo fastvideo-lora-v2 adapter, found {adapter_format!r}"
        )

    base_model = (metadata or {}).get("base_model")
    if base_model is not None and base_model != "MiniMaxAI/MiniMax-H3":
        raise ValueError(f"FastH3 adapter targets {base_model!r}, not MiniMaxAI/MiniMax-H3")

    replacement_count = int((metadata or {}).get("set_weight_tensors", "0"))
    vsa_modules = [module for module, suffixes in payloads.items() if _SET_WEIGHT in suffixes]
    if replacement_count or vsa_modules:
        raise ValueError(
            "This is a VSA FastH3 adapter. Native ComfyUI has no VSA-H3 attention backend or "
            "to_gate_compress parameters. Use the dense-datafree adapter with this node, or run "
            "the VSA adapter through FastVideo."
        )


def _build_key_map(lora, payloads, model_state):
    key_map = {}
    for source_module, suffixes in payloads.items():
        target = map_fastvideo_module(source_module)
        target_weight = f"diffusion_model.{target.module}.weight"
        if target_weight not in model_state:
            raise ValueError(
                f"FastH3 target {target_weight} is absent from the loaded model. "
                "Use a non-pruned MiniMax-H3 FL2VA base checkpoint."
            )

        if _DIFF_B in suffixes:
            if target.qkv_slot is not None:
                raise ValueError(f"Fused QKV bias deltas are unsupported for {source_module}")
            target_bias = f"diffusion_model.{target.module}.bias"
            if target_bias not in model_state:
                raise ValueError(f"FastH3 target {target_bias} is absent from the loaded model")

        if target.qkv_slot is None:
            key_map[source_module] = target_weight
            continue

        if _DIFF_B in suffixes or _SET_WEIGHT in suffixes:
            raise ValueError(f"Unsupported fused QKV payload for {source_module}")
        qkv_rows = int(model_state[target_weight].shape[0])
        if qkv_rows % 3:
            raise ValueError(f"ComfyUI QKV target has an invalid first dimension: {target_weight}")
        width = qkv_rows // 3
        if _LORA_B in suffixes and int(lora[source_module + _LORA_B].shape[0]) != width:
            raise ValueError(f"FastH3 LoRA output width does not match {target_weight}")
        key_map[source_module] = (target_weight, (0, target.qkv_slot * width, width))
    return key_map


class FastH3LoRALoader:
    """Load the dense FastVideo FastH3 Preview v1 adapter into native MiniMax-H3."""

    def __init__(self):
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "lora_name": (folder_paths.get_filename_list("loras"),),
                "strength_model": (
                    "FLOAT",
                    {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01},
                ),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_lora_model_only"
    CATEGORY = "model/loaders"
    DESCRIPTION = (
        "Loads FastVideo FastH3 Preview v1 dense adapters into native ComfyUI MiniMax-H3. "
        "VSA adapters and pruned AdaLN base checkpoints are intentionally rejected."
    )

    def _load_file(self, lora_path):
        if self.loaded_lora is not None and self.loaded_lora[0] == lora_path:
            return self.loaded_lora[1], self.loaded_lora[2]

        self.loaded_lora = None
        lora, metadata = comfy.utils.load_torch_file(
            lora_path, safe_load=True, return_metadata=True
        )
        self.loaded_lora = (lora_path, lora, metadata)
        return lora, metadata

    def load_lora_model_only(self, model, lora_name, strength_model):
        if strength_model == 0:
            return (model,)
        if not isinstance(model.model, comfy.model_base.MiniMaxH3):
            raise ValueError("FastH3 LoRA Loader requires a native ComfyUI MiniMax-H3 model")

        lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
        lora, metadata = self._load_file(lora_path)
        payloads, expected_patches = _analyze_payload(lora.keys())
        _validate_adapter(metadata, payloads)

        model_state = model.model.state_dict()
        if "diffusion_model.adaln_t_table" in model_state:
            raise ValueError(
                "FastH3 Preview v1 targets full-width AdaLN weights. The loaded MiniMax-H3 base "
                "is a pruned/curve-basis checkpoint; load a non-pruned FL2VA checkpoint instead."
            )

        key_map = _build_key_map(lora, payloads, model_state)
        patches = comfy.lora.load_lora(lora, key_map, log_missing=False)
        if len(patches) != expected_patches:
            raise RuntimeError(
                f"FastH3 adapter conversion produced {len(patches)} patches; "
                f"{expected_patches} were expected"
            )

        patched_model = model.clone()
        applied = patched_model.add_patches(patches, strength_model)
        if len(applied) != len(patches):
            raise RuntimeError(
                f"ComfyUI accepted {len(applied)} of {len(patches)} FastH3 patches"
            )
        if metadata:
            patched_model.set_attachments("lora_metadata", metadata)

        logging.info(
            "Loaded FastH3 dense adapter %s at strength %g (%d patches)",
            os.path.basename(lora_path),
            strength_model,
            len(patches),
        )
        return (patched_model,)


NODE_CLASS_MAPPINGS = {"FastH3LoRALoader": FastH3LoRALoader}
NODE_DISPLAY_NAME_MAPPINGS = {
    "FastH3LoRALoader": "FastH3 LoRA Loader (Dense)",
}


# mapping.py

"""Name translation for FastVideo MiniMax-H3 adapters.

This module deliberately has no ComfyUI or torch imports so its mapping rules can
be tested without loading either runtime.
"""

from dataclasses import dataclass
import re


PAYLOAD_SUFFIXES = (
    ".lora_A.weight",
    ".lora_B.weight",
    ".set_weight",
    ".diff_b",
    ".diff",
    ".alpha",
)


@dataclass(frozen=True)
class ModuleTarget:
    module: str
    qkv_slot: int | None = None


def split_payload_key(key: str) -> tuple[str, str]:
    for suffix in PAYLOAD_SUFFIXES:
        if key.endswith(suffix):
            return key[: -len(suffix)], suffix
    raise ValueError(f"unsupported FastH3 adapter tensor: {key}")


def _map_block_rest(rest: str) -> ModuleTarget:
    qkv = re.fullmatch(r"attn\.to_([qkv])", rest)
    if qkv:
        return ModuleTarget("attn.qkv_proj", {"q": 0, "k": 1, "v": 2}[qkv.group(1)])

    direct = {
        "attn.to_out.0": "attn.out_proj",
        "ff.net.0.proj": "mlp.fc1",
        "ff.net.2": "mlp.fc2",
        "adaln_proj.linear": "adaln_proj.linear",
        "norm1": "norm1",
        "norm2": "norm2",
    }
    target = direct.get(rest)
    if target is None:
        raise ValueError(f"unsupported FastVideo MiniMax-H3 block module: {rest}")
    return ModuleTarget(target)


def map_fastvideo_module(module: str) -> ModuleTarget:
    boundary = {
        "proj_in": "video_patch_proj",
        "audio_proj_in": "audio_patch_proj",
        "context_embedder": "condition_proj",
        "time_embedder.linear_1": "time_embedder.proj_in",
        "time_embedder.linear_2": "time_embedder.proj_out",
        "norm_out.norm": "final_layer.norm",
        "norm_out.linear": "final_layer.adaln_proj.linear",
        "proj_out": "final_layer.video_out",
        "audio_proj_out": "final_layer.audio_out",
    }
    target = boundary.get(module)
    if target is not None:
        return ModuleTarget(target)

    block = re.fullmatch(r"transformer_blocks\.(\d+)\.(.+)", module)
    if block:
        mapped = _map_block_rest(block.group(2))
        return ModuleTarget(f"blocks.{block.group(1)}.{mapped.module}", mapped.qkv_slot)

    refiner = re.fullmatch(r"token_refiner\.refiner_blocks\.(\d+)\.(.+)", module)
    if refiner:
        mapped = _map_block_rest(refiner.group(2))
        if mapped.module.startswith("adaln_proj."):
            raise ValueError("MiniMax-H3 token-refiner blocks do not have AdaLN projections")
        return ModuleTarget(f"token_refiner.blocks.{refiner.group(1)}.{mapped.module}", mapped.qkv_slot)

    raise ValueError(f"unsupported FastVideo MiniMax-H3 module: {module}")


# __init__.py

from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]