"""One ComfyUI loader for recovered MiniMax-H3 text encoders.

This extension adds one visible node. It loads a selected text-only encoder package
and returns a standard CLIP object for official MiniMax-H3 nodes.
"""
from __future__ import annotations

import hashlib
import json
import logging
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from safetensors.torch import load_file

import comfy.sd
import comfy.sd1_clip
import comfy.supported_models_base
import comfy.utils
import comfy.model_management
import comfy.weight_adapter
import folder_paths
from comfy.text_encoders.hunyuan_video import llama_detect
from comfy.text_encoders.minimax import MiniMaxQwen3VL
from comfy.text_encoders.qwen3vl import Qwen3VLSDTokenizer

MODEL_ROOT = Path(folder_paths.models_dir) / "text_encoders"
RECOVERED_FORMAT = "minimax_h3_recovered_8b_release_v1"
PRUNED24_FORMAT = "minimax_h3_pruned24_balanced_release_v1"
ARA_TARGETS = ("self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj")


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def _load_manifest(directory: Path, filename: str) -> dict[str, Any]:
    path = directory / filename
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"Invalid manifest: {path}")
    return value


def _verified(directory: Path, entry: dict[str, Any]) -> Path:
    path = (directory / entry["file"]).resolve()
    if path.parent != directory.resolve():
        raise ValueError(f"Model file must be inside its package: {path}")
    actual = _sha256(path)
    if actual != entry["sha256"]:
        raise ValueError(f"SHA-256 mismatch for {path.name}: {actual} != {entry['sha256']}")
    return path


def _submodule(root: Any, dotted_name: str) -> Any:
    current = root
    for part in dotted_name.split("."):
        current = current[int(part)] if part.isdigit() else getattr(current, part)
    return current


def _set_submodule(root: Any, dotted_name: str, value: nn.Module) -> None:
    parts = dotted_name.split(".")
    parent = root
    for part in parts[:-1]:
        parent = parent[int(part)] if part.isdigit() else getattr(parent, part)
    if parts[-1].isdigit():
        parent[int(parts[-1])] = value
    else:
        setattr(parent, parts[-1], value)


class IdentityBlock(nn.Module):
    def forward(self, x, *args, **kwargs):
        return x, None


class RecoveryLinear(nn.Module):
    def __init__(self, base: nn.Module, rank: int, alpha: float, device: torch.device, prefix: str = "lora"):
        super().__init__()
        self.base = base
        self.prefix = prefix
        setattr(self, f"{prefix}_down", nn.Linear(base.in_features, rank, bias=False, device=device, dtype=torch.float32))
        setattr(self, f"{prefix}_up", nn.Linear(rank, base.out_features, bias=False, device=device, dtype=torch.float32))
        self.scaling = alpha / rank

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base_output = self.base(x)
        down = getattr(self, f"{self.prefix}_down")
        up = getattr(self, f"{self.prefix}_up")
        delta = up(down(x.float())) * self.scaling
        return base_output + delta.to(dtype=base_output.dtype)


class ConditioningAdapter(nn.Module):
    def __init__(self, input_width: int, output_width: int, bottleneck: int, device: torch.device):
        super().__init__()
        self.norm = nn.RMSNorm(input_width, device=device, dtype=torch.float32)
        self.proj = nn.Linear(input_width, output_width, bias=False, device=device, dtype=torch.float32)
        self.down = nn.Linear(input_width, bottleneck, bias=False, device=device, dtype=torch.float32)
        self.up = nn.Linear(bottleneck, output_width, bias=False, device=device, dtype=torch.float32)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        normalized = self.norm(hidden.float())
        return self.proj(normalized) + self.up(torch.nn.functional.silu(self.down(normalized)))


class RecoveredMiniMaxTextModel(MiniMaxQwen3VL):
    model_type = "qwen3vl_8b"

    def __init__(self, config_dict, dtype, device, operations):
        super().__init__({"num_hidden_layers": 24, "final_norm": False, "lm_head": False}, dtype, device, operations)
        # ponytail: text-only package; tokenizer rejects image/reference inputs.
        del self.visual
        adapter = config_dict["adapter"]
        for layer in config_dict["ara_layers"]:
            for suffix in ARA_TARGETS:
                parent_name, leaf = f"layers.{layer}.{suffix}".rsplit(".", 1)
                parent = _submodule(self.model, parent_name)
                base = getattr(parent, leaf)
                setattr(parent, leaf, RecoveryLinear(base, config_dict["ara_rank"], config_dict["ara_alpha"], device))
        self.adapter = ConditioningAdapter(adapter["input_width"], adapter["output_width"], adapter["bottleneck"], device)

    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None,
                final_layer_norm_intermediate=True, dtype=None, embeds_info=None, **kwargs):
        result = super().forward(
            input_ids, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens,
            intermediate_output=None, final_layer_norm_intermediate=False, dtype=dtype,
            embeds_info=embeds_info or [], **kwargs,
        )
        return self.adapter(result[0]), None


def _recovered_clip_model(recovery: dict[str, Any], quantization_metadata: dict[str, Any] | None):
    class RecoveredClipModel(comfy.sd1_clip.SDClipModel):
        def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options=None):
            options = dict(model_options or {})
            if quantization_metadata is not None:
                options["quantization_metadata"] = quantization_metadata
            super().__init__(
                device=device, layer="last", layer_idx=None, textmodel_json_config={
                    "adapter": recovery["adapter"], "ara_layers": recovery["ara"]["layers"],
                    "ara_rank": recovery["ara"]["rank"], "ara_alpha": recovery["ara"]["alpha"],
                }, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False,
                model_class=RecoveredMiniMaxTextModel, enable_attention_masks=False,
                return_attention_masks=False, model_options=options,
            )

        def encode_token_weights(self, token_weight_pairs):
            output = super().encode_token_weights(token_weight_pairs)
            tags = getattr(self.transformer, "last_token_tags", None)
            if tags is None:
                tags = torch.ones(output[0].shape[1], dtype=torch.long)
            extra = output[2] if len(output) > 2 and isinstance(output[2], dict) else {}
            extra["minimax_token_tags"] = tags
            return output[0], output[1], extra
    return RecoveredClipModel


def _recovered_te_model(recovery: dict[str, Any], quantization_metadata: dict[str, Any] | None):
    clip_model = _recovered_clip_model(recovery, quantization_metadata)
    class RecoveredTEModel(comfy.sd1_clip.SD1ClipModel):
        def __init__(self, device="cpu", dtype=None, model_options=None):
            super().__init__(device=device, dtype=dtype, name="qwen3vl_8b", clip_model=clip_model, model_options=model_options or {})
    return RecoveredTEModel


class RecoveredMiniMaxTokenizer(comfy.sd1_clip.SD1Tokenizer):
    """The official H3 verbatim text token contract, without vision inputs."""
    def __init__(self, embedding_directory=None, tokenizer_data=None):
        tokenizer = lambda *args, **kwargs: Qwen3VLSDTokenizer(*args, **kwargs, embedding_size=4096, embedding_key="qwen3vl_8b")
        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data or {}, name="qwen3vl_8b", tokenizer=tokenizer)

    def tokenize_with_weights(self, text, return_word_ids=False, images=None, minimax_ref_items=None, **kwargs):
        if images or minimax_ref_items:
            raise ValueError("Recovered MiniMax-H3 encoders support text-only T2V. Do not use image or reference inputs.")
        tokens = [(token, 1.0) for token in self.qwen3vl_8b.tokenizer(text, add_special_tokens=False)["input_ids"]]
        if not tokens:
            tokens = [(151643, 1.0)]
        if return_word_ids:
            tokens = [token + (0,) for token in tokens]
        return {"qwen3vl_8b": [tokens]}


def _translate_recovered_base(state: dict[str, torch.Tensor], ara_layers: list[int]) -> dict[str, torch.Tensor]:
    translated: dict[str, torch.Tensor] = {}
    for key, value in state.items():
        output = key
        for layer in ara_layers:
            for suffix in ARA_TARGETS:
                prefix = f"model.layers.{layer}.{suffix}."
                if key.startswith(prefix):
                    output = prefix + "base." + key.removeprefix(prefix)
                    break
        translated[output] = value
    return translated


def _load_recovered(directory: Path, package: dict[str, Any], variant_id: str):
    recovery = {
        "ara": package["ara"],
        "adapter": package["adapter"],
    }
    variant = package["base_variants"][variant_id]
    base_path = _verified(directory, variant)
    adapter_path = _verified(directory, recovery["adapter"])
    ara_path = _verified(directory, recovery["ara"])
    state, metadata = comfy.utils.load_torch_file(str(base_path), safe_load=True, return_metadata=True)
    state, _ = comfy.utils.convert_old_quants(state, metadata=metadata)
    quantization_metadata = llama_detect(state)["llama_quantization_metadata"] if any(
        key.endswith(".comfy_quant") for key in state
    ) else None
    target = comfy.supported_models_base.ClipTarget(
        RecoveredMiniMaxTokenizer, _recovered_te_model(recovery, quantization_metadata)
    )
    # The base, ARA, and adapter state is loaded below. Keep the empty model on
    # CPU during construction; otherwise CLIP forces GPU placement before the
    # custom RecoveryLinear wrappers have populated weights.
    clip = comfy.sd.CLIP(
        target,
        embedding_directory=folder_paths.get_folder_paths("embeddings"),
        parameters=6_000_000_000,
        model_options={"initial_device": comfy.model_management.text_encoder_offload_device()},
    )
    state = _translate_recovered_base(state, recovery["ara"]["layers"])
    adapter = {f"adapter.{key}": value for key, value in load_file(str(adapter_path), device="cpu").items()}
    ara = {f"model.{key}": value for key, value in load_file(str(ara_path), device="cpu").items()}
    state.update(adapter)
    state.update(ara)
    missing, unexpected = clip.load_sd(state, full_model=False)
    real_missing = [key for key in (missing or []) if not key.endswith("lm_head.weight")]
    if real_missing or unexpected:
        raise RuntimeError(f"Recovered encoder load mismatch: missing={real_missing[:8]} unexpected={list(unexpected or [])[:8]}")
    return clip


def _load_pruned24(directory: Path, package: dict[str, Any], variant_id: str):
    checkpoint = _verified(directory, package["base_variants"][variant_id])
    ara_path = _verified(directory, package["ara"])
    clip = comfy.sd.load_clip([str(checkpoint)], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=comfy.sd.CLIPType.MINIMAX)
    transformer = clip.cond_stage_model.qwen3vl_32b.transformer
    for layer in package["removed_blocks"]:
        transformer.model.layers[layer] = IdentityBlock()
    ara = load_file(str(ara_path), device="cpu")
    rank = package["ara"]["rank"]
    alpha = package["ara"]["alpha"]
    expected = set()
    manager = comfy.weight_adapter.BypassInjectionManager()
    for layer in package["ara"]["layers"]:
        for suffix in ARA_TARGETS:
            prefix = f"layers.{layer}.{suffix}"
            down_key = f"{prefix}.ara_down.weight"
            up_key = f"{prefix}.ara_up.weight"
            adapter = comfy.weight_adapter.LoRAAdapter(
                {down_key, up_key},
                (ara[up_key], ara[down_key], alpha, None, None, None),
            )
            manager.add_adapter(f"qwen3vl_32b.transformer.model.{prefix}", adapter, strength=1.0)
            expected.update({down_key, up_key})
    if set(ara) != expected:
        raise RuntimeError("Pruned-24 ARA keys do not match manifest")
    injections = manager.create_injections(clip.cond_stage_model)
    if manager.get_hook_count() != len(expected) // 2:
        raise RuntimeError("Pruned-24 ARA injection count does not match manifest")
    clip.patcher.set_injections("balanced_ara", injections)
    clip.patcher.cached_patcher_init = (_reload_pruned24, (directory, package, variant_id))
    return clip


def _reload_pruned24(directory: Path, package: dict[str, Any], variant_id: str):
    return _load_pruned24(directory, package, variant_id).patcher


def _choices() -> dict[str, tuple[Path, dict[str, Any], str | None]]:
    choices: dict[str, tuple[Path, dict[str, Any], str | None]] = {}
    if not MODEL_ROOT.is_dir():
        return choices
    for directory in MODEL_ROOT.iterdir():
        recovered = directory / "minimax_h3_recovered_8b_manifest.json"
        pruned = directory / "minimax_h3_pruned24_balanced_manifest.json"
        if recovered.is_file():
            package = _load_manifest(directory, recovered.name)
            if package.get("format") == RECOVERED_FORMAT:
                for variant in sorted(package["base_variants"]):
                    choices[f"{directory.name}: {variant}"] = (directory, package, variant)
        if pruned.is_file():
            package = _load_manifest(directory, pruned.name)
            if package.get("format") == PRUNED24_FORMAT:
                for variant_id in sorted(package["base_variants"]):
                    choices[f"{directory.name}: {variant_id}"] = (directory, package, variant_id)
    return choices


class MiniMaxH3TextEncoderLoader:
    @classmethod
    def INPUT_TYPES(cls):
        choices = sorted(_choices())
        return {"required": {"text_encoder": (choices,)}}

    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load"
    CATEGORY = "advanced/loaders"
    DESCRIPTION = "Loads one recovered MiniMax-H3 text encoder. Text-only T2V."

    def load(self, text_encoder):
        choice = _choices().get(text_encoder)
        if choice is None:
            raise ValueError("Selected package is missing. Reopen the workflow after model files are installed.")
        directory, package, variant = choice
        if package["format"] == RECOVERED_FORMAT:
            clip = _load_recovered(directory, package, variant)
        elif package["format"] == PRUNED24_FORMAT:
            clip = _load_pruned24(directory, package, variant)
        else:
            raise ValueError("Unsupported MiniMax-H3 text encoder package")
        logging.info("MiniMaxH3TextEncoderLoader: %s", text_encoder)
        return (clip,)


NODE_CLASS_MAPPINGS = {"MiniMaxH3TextEncoderLoader": MiniMaxH3TextEncoderLoader}
NODE_DISPLAY_NAME_MAPPINGS = {"MiniMaxH3TextEncoderLoader": "Load MiniMax H3 Text Encoder"}