"""Cosmos-Reference (Mirumo0u0) for DiffSynth Anima — Anima Edit support.

Implements the same inference contract as
https://github.com/Mirumo0u0/ComfyUI-Cosmos-Reference :

  * VAE-encode a source image to a reference latent
  * concatenate the reference latent as an extra temporal frame on every DiT step
  * slice the reference frame(s) off the DiT output before the sampler advances
  * no attention bias / strength gating (unlike In-Context Character); control
    strength comes from the Edit LoRA and the denoising strength

This is the mechanism required by Anima Edit v1 / v2 and other edit LoRAs that
expect Cosmos-style reference latents.
"""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterator, Optional

import torch
from PIL import Image

# Official ComfyUI-Cosmos-Reference constants (kept for clarity / debugging).
MAX_REF_LATENTS = 1
COND_REF_LATENTS_KEY = "ref_latents"


@dataclass
class CosmosReferenceState:
    """Request-scoped state for one Cosmos-Reference generation."""

    ref_latents: list[torch.Tensor] = field(default_factory=list)  # each (C, H, W)
    enabled: bool = False
    active: bool = False


_STATE = CosmosReferenceState()


def get_state() -> CosmosReferenceState:
    return _STATE


def clear_state() -> None:
    _STATE.ref_latents = []
    _STATE.enabled = False
    _STATE.active = False


def encode_reference_image(
    pipe: Any,
    image: Image.Image | None,
    width: int,
    height: int,
) -> Optional[torch.Tensor]:
    """Encode one RGB image with the Anima VAE to a single-frame latent (C, h, w).

    Matches the resize + centre-pad behaviour used by typical Anima Edit
    workflows (letterbox into the generation canvas, then VAE-encode).
    """
    if image is None:
        return None
    if not isinstance(image, Image.Image):
        try:
            image = Image.fromarray(image)
        except Exception:
            return None
    image = image.convert("RGB")

    src_w, src_h = image.size
    if src_w < 1 or src_h < 1:
        return None

    # Letterbox into target canvas (same idea as the sample workflow resize node).
    scale = min(width / src_w, height / src_h)
    new_w = max(1, int(round(src_w * scale)))
    new_h = max(1, int(round(src_h * scale)))
    resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (width, height), (0, 0, 0))
    canvas.paste(resized, ((width - new_w) // 2, (height - new_h) // 2))

    device = pipe.device
    dtype = pipe.torch_dtype
    vae = pipe.vae

    if hasattr(pipe, "load_models_to_device"):
        pipe.load_models_to_device(["vae"])

    with torch.no_grad():
        if hasattr(pipe, "preprocess_image"):
            tensor = pipe.preprocess_image(canvas).to(device=device, dtype=dtype)
        else:
            import numpy as np

            arr = np.asarray(canvas, dtype="float32") / 127.5 - 1.0
            tensor = (
                torch.from_numpy(arr)
                .to(device=device, dtype=dtype)
                .permute(2, 0, 1)
            )

        # Match DiffSynth / Anima VAE encode paths used elsewhere in the app.
        try:
            if tensor.ndim == 3:
                enc = vae.encode(tensor.unsqueeze(2), device=device)
            else:
                enc = vae.encode(tensor.unsqueeze(2), device=device)
        except TypeError:
            video = tensor if tensor.ndim == 3 else tensor[0]
            if video.ndim == 3:
                video = video.unsqueeze(1)
            enc = vae.encode([video], device=device)
        except Exception:
            video = tensor if tensor.ndim == 3 else tensor[0]
            if video.ndim == 3:
                video = video.unsqueeze(1)
            enc = vae.encode([video], device=device)

        if isinstance(enc, (list, tuple)):
            enc = torch.stack(
                [e if isinstance(e, torch.Tensor) else e for e in enc]
            )
        if enc.ndim == 5:
            frame = enc[0, :, 0] if enc.shape[0] >= 1 else enc.squeeze(0)[:, 0]
        elif enc.ndim == 4:
            if enc.shape[0] == 1:
                frame = enc[0]
            else:
                frame = enc[:, 0] if enc.shape[1] < enc.shape[0] else enc[0]
        elif enc.ndim == 3:
            frame = enc
        else:
            raise RuntimeError(f"Unexpected VAE encode shape: {tuple(enc.shape)}")

        return frame.detach().to(device="cpu", dtype=torch.bfloat16).contiguous()


def cosmos_model_fn(
    dit=None,
    latents=None,
    timestep=None,
    prompt_emb=None,
    t5xxl_ids=None,
    t5xxl_weights=None,
    use_gradient_checkpointing=False,
    use_gradient_checkpointing_offload=False,
    **kwargs,
):
    """DiT forward with optional Cosmos-Reference temporal concatenation.

    Matches the official ``cosmos_diffusion_reference_wrapper`` exactly:
    * concat ref frames on the temporal axis of ``x`` only
    * do **not** rewrite timesteps (Comfy leaves the original ``t`` intact)
    * slice the output back to the original temporal length
    """
    state = get_state()
    # Pipeline passes (B, C, H, W); DiT expects a temporal axis.
    if latents.ndim == 4:
        x = latents.unsqueeze(2)  # (B, C, 1, H, W)
        squeezed = True
    else:
        x = latents
        squeezed = False
    t = timestep / 1000 if timestep is not None else timestep

    use_ref = state.enabled and bool(state.ref_latents)
    if not use_ref:
        model_output = dit(
            x=x,
            timesteps=t,
            context=prompt_emb,
            t5xxl_ids=t5xxl_ids,
            t5xxl_weights=t5xxl_weights,
            use_gradient_checkpointing=use_gradient_checkpointing,
            use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
        )
        if squeezed and getattr(model_output, "ndim", 0) == 5:
            return model_output.squeeze(2)
        return model_output

    # Official node: only extend x; leave timesteps unchanged.
    x_temporal_dim = int(x.shape[2])
    B, C, _, H, W = x.shape
    pieces = [x]
    for ref in state.ref_latents:
        r = ref.to(device=x.device, dtype=x.dtype)
        if r.ndim == 3:  # (C, H, W)
            r = r.unsqueeze(0).unsqueeze(2)  # (1, C, 1, H, W)
        elif r.ndim == 4:  # (B|1, C, H, W)
            r = r.unsqueeze(2)
        if r.shape[0] == 1 and B > 1:
            r = r.expand(B, -1, -1, -1, -1)
        if r.shape[-2] != H or r.shape[-1] != W:
            b, c, tt = r.shape[0], r.shape[1], r.shape[2]
            r = r.reshape(b * tt, c, r.shape[-2], r.shape[-1])
            r = torch.nn.functional.interpolate(
                r, size=(H, W), mode="bilinear", align_corners=False
            )
            r = r.reshape(b, c, tt, H, W)
        pieces.append(r)
    x_cat = torch.cat(pieces, dim=2)

    state.active = True
    try:
        model_output = dit(
            x=x_cat,
            timesteps=t,  # unchanged — matches ComfyUI-Cosmos-Reference
            context=prompt_emb,
            t5xxl_ids=t5xxl_ids,
            t5xxl_weights=t5xxl_weights,
            use_gradient_checkpointing=use_gradient_checkpointing,
            use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
        )
    finally:
        state.active = False

    if getattr(model_output, "ndim", 0) == 5:
        model_output = model_output[:, :, :x_temporal_dim]
        if squeezed:
            model_output = model_output.squeeze(2)
    return model_output


@contextmanager
def temporary_cosmos_reference(
    pipe: Any,
    source_image: Image.Image | None,
    width: int,
    height: int,
) -> Iterator[bool]:
    """Encode the source image and enable Cosmos-Reference for the duration of the block."""
    clear_state()
    if source_image is None:
        yield False
        return
    try:
        latent = encode_reference_image(pipe, source_image, width, height)
        if latent is None:
            yield False
            return
        st = get_state()
        st.ref_latents = [latent]
        st.enabled = True
        yield True
    finally:
        clear_state()