"""ComfyUI nodes for the Anima style stream (decoupled cross-attention).

Requires a trained AnimaStyleAdapter checkpoint (see style_adapter.py
and the project docs); with an untrained adapter the gates are zero and
the nodes are an exact no-op.

The style stream composes freely with the in-context character stream:
AnimaInContextApply patches self-attention (T-axis reference frames),
AnimaStyleApply patches cross-attention — chain both Apply nodes and
control char/style strength independently.
"""

import os

import torch

import comfy.utils
import folder_paths
from comfy.patcher_extension import WrappersMP

from .style_adapter import AnimaStyleAdapter, StyleState

STYLE_WRAPPER_KEY = "anima_style_ref"

_ADAPTER_DIR = "anima_style_adapters"
if _ADAPTER_DIR not in folder_paths.folder_names_and_paths:
    folder_paths.add_model_folder_path(
        _ADAPTER_DIR, os.path.join(folder_paths.models_dir, _ADAPTER_DIR)
    )


class AnimaStyleAdapterLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "adapter_name": (folder_paths.get_filename_list(_ADAPTER_DIR),),
            },
        }

    RETURN_TYPES = ("ANIMA_STYLE_ADAPTER",)
    FUNCTION = "load"
    CATEGORY = "anima/style"

    def load(self, adapter_name):
        path = folder_paths.get_full_path_or_raise(_ADAPTER_DIR, adapter_name)
        sd = comfy.utils.load_torch_file(path, safe_load=True)
        adapter = AnimaStyleAdapter.from_state_dict(sd)
        adapter.eval()
        return (adapter,)


class AnimaStyleEncode:
    """Encode style reference image(s) with a SigLIP CLIP_VISION model,
    keeping the hidden-state stack the adapter aggregates over.

    Multiple images are encoded independently and their patch tokens are
    concatenated at apply time (attention pools over all of them)."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip_vision": ("CLIP_VISION",),
                "image": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("ANIMA_STYLE_EMBEDS",)
    FUNCTION = "encode"
    CATEGORY = "anima/style"

    def encode(self, clip_vision, image):
        out = clip_vision.encode_image(image)
        hs = out["all_hidden_states"]  # (B, L, N, D)
        if hs is None:
            raise RuntimeError(
                "this CLIP_VISION model does not expose all hidden states; "
                "use a SigLIP vision model"
            )
        return ({"hidden_states": hs},)


class AnimaStyleApply:
    """Attach the style stream to an Anima model.

    style_weight: global multiplier on the (learned, per-block) gates.
    cond_only:    apply style only to the cond half of the CFG batch.
    start/end_percent: sampling window, same semantics as the
    in-context Apply node.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "style_adapter": ("ANIMA_STYLE_ADAPTER",),
                "style_embeds": ("ANIMA_STYLE_EMBEDS",),
                "style_weight": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 5.0, "step": 0.05}),
                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
            },
            "optional": {
                "cond_only": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "apply"
    CATEGORY = "anima/style"

    def apply(self, model, style_adapter, style_embeds, style_weight,
              start_percent, end_percent, cond_only=True):
        m = model.clone()

        ms = m.get_model_object("model_sampling")
        sigma_start = ms.percent_to_sigma(start_percent)
        sigma_end = ms.percent_to_sigma(end_percent)

        dm = m.get_model_object("diffusion_model")
        n_layers = style_adapter.config["n_layers"]
        hidden_states = style_embeds["hidden_states"][:, -n_layers:]  # (B, K, N, D)

        state = StyleState()
        # per-(device, dtype) cache of the per-block K/V — the style
        # tokens are constant across sampling steps
        kv_cache = {}

        def compute_kv(device, dtype):
            key = (device, dtype)
            if key not in kv_cache:
                adapter = style_adapter.to(device)
                hs = hidden_states.to(device=device, dtype=next(adapter.parameters()).dtype)
                with torch.no_grad():
                    tokens = adapter.aggregator(hs)  # (B, N, style_dim)
                    # multiple style images -> one long token sequence
                    tokens = tokens.reshape(1, -1, tokens.shape[-1])
                    kv_cache[key] = [
                        tuple(t.to(dtype) for t in blk.kv(tokens)) for blk in adapter.blocks
                    ]
            return kv_cache[key]

        def wrapper(executor, x, timesteps, context, fps=None, padding_mask=None, **kwargs):
            to = kwargs.get("transformer_options", {})
            sigmas = to.get("sigmas", None)
            if sigmas is not None:
                s = float(sigmas.max())
                if s > sigma_start or s < sigma_end:
                    return executor(x, timesteps, context, fps, padding_mask, **kwargs)

            state.kv_per_block = compute_kv(x.device, x.dtype)
            state.weight = style_weight

            state.sample_scale_B = None
            cou = to.get("cond_or_uncond", None)
            if cond_only and cou and x.shape[0] % len(cou) == 0:
                chunk = x.shape[0] // len(cou)
                scale = torch.ones(x.shape[0], device=x.device)
                for i, kind in enumerate(cou):
                    if kind == 1:
                        scale[i * chunk:(i + 1) * chunk] = 0.0
                state.sample_scale_B = scale

            state.active = True
            try:
                return executor(x, timesteps, context, fps, padding_mask, **kwargs)
            finally:
                state.active = False

        m.add_wrapper_with_key(WrappersMP.DIFFUSION_MODEL, STYLE_WRAPPER_KEY, wrapper)

        from .style_adapter import StyleCrossAttention

        for i, block in enumerate(dm.blocks):
            m.add_object_patch(
                "diffusion_model.blocks.{}.cross_attn".format(i),
                StyleCrossAttention(block.cross_attn, style_adapter.blocks[i], state, i),
            )

        return (m,)


NODE_CLASS_MAPPINGS = {
    "AnimaStyleAdapterLoader": AnimaStyleAdapterLoader,
    "AnimaStyleEncode": AnimaStyleEncode,
    "AnimaStyleApply": AnimaStyleApply,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AnimaStyleAdapterLoader": "Anima Style Adapter Loader",
    "AnimaStyleEncode": "Anima Style Encode (SigLIP)",
    "AnimaStyleApply": "Anima Style Apply",
}