"""VAE loading and decode support for Z-Image Turbo."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F

from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import (
    LATENT_CHANNELS,
    LATENT_DOWNSCALE,
    VAE_ROLE,
    ModelSource,
)
from diffusion_cli.errors import DiffusionCliError


FLUX_SCALE_FACTOR = 0.3611
FLUX_SHIFT_FACTOR = 0.1159
VAE_CHANNELS = 128
VAE_CHANNEL_MULTIPLIERS = (1, 2, 4, 4)
VAE_NUM_RES_BLOCKS = 2
VAE_OUT_CHANNELS = 3
VAE_RESOLUTION = 256


@dataclass(frozen=True)
class VaeConfig:
    """Architecture settings for the supported image VAE decoder."""

    channels: int = VAE_CHANNELS
    channel_multipliers: tuple[int, ...] = VAE_CHANNEL_MULTIPLIERS
    num_res_blocks: int = VAE_NUM_RES_BLOCKS
    latent_channels: int = LATENT_CHANNELS
    out_channels: int = VAE_OUT_CHANNELS
    resolution: int = VAE_RESOLUTION


def fluxLatentToVae(latent: torch.Tensor) -> torch.Tensor:
    """Convert Flux-format model latents into VAE latents."""

    return latent / FLUX_SCALE_FACTOR + FLUX_SHIFT_FACTOR


def postprocessVaeOutput(image: torch.Tensor) -> torch.Tensor:
    """Convert decoded VAE output from [-1, 1] to clamped [0, 1]."""

    return image.add(1.0).div(2.0).clamp(0.0, 1.0)


def detectVaeConfig(state_dict: dict[str, torch.Tensor]) -> VaeConfig:
    """Detect the exact Z-Image VAE decoder architecture."""

    try:
        conv_in = state_dict["decoder.conv_in.weight"]
        conv_out = state_dict["decoder.conv_out.weight"]
    except KeyError as exc:
        raise DiffusionCliError(
            "Unsupported VAE state dict: no known decoder keys found"
        ) from exc

    if conv_in.ndim != 4 or conv_in.shape[1] != LATENT_CHANNELS:
        raise DiffusionCliError(
            "Unsupported VAE latent channels: expected "
            f"{LATENT_CHANNELS}, got {tuple(conv_in.shape)}"
        )
    if conv_out.ndim != 4 or conv_out.shape[0] != VAE_OUT_CHANNELS:
        raise DiffusionCliError(
            "Unsupported VAE output channels: expected "
            f"{VAE_OUT_CHANNELS}, got {tuple(conv_out.shape)}"
        )

    base_channels = conv_out.shape[1]
    block_in = conv_in.shape[0]
    if base_channels != VAE_CHANNELS or block_in != VAE_CHANNELS * 4:
        raise DiffusionCliError(
            "Unsupported VAE channel layout: expected Z-Image image VAE"
        )

    return VaeConfig()


def _normalize(channels: int) -> nn.GroupNorm:
    return nn.GroupNorm(
        num_groups=32,
        num_channels=channels,
        eps=1e-6,
        affine=True,
    )


class Upsample(nn.Module):
    """Nearest-neighbor 2x upsample followed by a 3x3 convolution."""

    def __init__(self, channels: int, *, dtype=None, device=None) -> None:
        super().__init__()
        self.conv = nn.Conv2d(
            channels,
            channels,
            kernel_size=3,
            padding=1,
            dtype=dtype,
            device=device,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply 2x nearest-neighbor upsampling."""

        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
        return self.conv(x)


class ResnetBlock(nn.Module):
    """Residual block used by the Z-Image VAE decoder."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int | None = None,
        *,
        dtype=None,
        device=None,
    ) -> None:
        super().__init__()
        out_channels = out_channels if out_channels is not None else in_channels
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.norm1 = _normalize(in_channels).to(device=device, dtype=dtype)
        self.conv1 = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=3,
            padding=1,
            dtype=dtype,
            device=device,
        )
        self.norm2 = _normalize(out_channels).to(device=device, dtype=dtype)
        self.dropout = nn.Dropout(0.0)
        self.conv2 = nn.Conv2d(
            out_channels,
            out_channels,
            kernel_size=3,
            padding=1,
            dtype=dtype,
            device=device,
        )
        if in_channels != out_channels:
            self.nin_shortcut = nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=1,
                dtype=dtype,
                device=device,
            )
        else:
            self.nin_shortcut = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply the residual update."""

        h = self.norm1(x)
        h = F.silu(h)
        h = self.conv1(h)
        h = self.norm2(h)
        h = F.silu(h)
        h = self.dropout(h)
        h = self.conv2(h)
        if self.nin_shortcut is not None:
            x = self.nin_shortcut(x)
        return x + h


class AttnBlock(nn.Module):
    """Single spatial self-attention block used at the VAE bottleneck."""

    def __init__(self, channels: int, *, dtype=None, device=None) -> None:
        super().__init__()
        self.norm = _normalize(channels).to(device=device, dtype=dtype)
        self.q = nn.Conv2d(
            channels,
            channels,
            kernel_size=1,
            dtype=dtype,
            device=device,
        )
        self.k = nn.Conv2d(
            channels,
            channels,
            kernel_size=1,
            dtype=dtype,
            device=device,
        )
        self.v = nn.Conv2d(
            channels,
            channels,
            kernel_size=1,
            dtype=dtype,
            device=device,
        )
        self.proj_out = nn.Conv2d(
            channels,
            channels,
            kernel_size=1,
            dtype=dtype,
            device=device,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply full spatial attention over latent positions."""

        batch_size, channels, height, width = x.shape
        h = self.norm(x)
        query = self.q(h).reshape(batch_size, channels, -1).transpose(1, 2)
        key = self.k(h).reshape(batch_size, channels, -1).transpose(1, 2)
        value = self.v(h).reshape(batch_size, channels, -1).transpose(1, 2)
        h = F.scaled_dot_product_attention(
            query.unsqueeze(1),
            key.unsqueeze(1),
            value.unsqueeze(1),
        )
        h = h.squeeze(1).transpose(1, 2).reshape(
            batch_size,
            channels,
            height,
            width,
        )
        return x + self.proj_out(h)


class VaeDecoder(nn.Module):
    """Image-only decoder matching ComfyUI's Z-Image VAE branch."""

    def __init__(
        self,
        config: VaeConfig,
        *,
        dtype=None,
        device=None,
    ) -> None:
        super().__init__()
        self.config = config
        self.num_resolutions = len(config.channel_multipliers)
        self.num_res_blocks = config.num_res_blocks
        block_in = (
            config.channels
            * config.channel_multipliers[self.num_resolutions - 1]
        )
        self.conv_in = nn.Conv2d(
            config.latent_channels,
            block_in,
            kernel_size=3,
            padding=1,
            dtype=dtype,
            device=device,
        )
        self.mid = nn.Module()
        self.mid.block_1 = ResnetBlock(
            block_in,
            block_in,
            dtype=dtype,
            device=device,
        )
        self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
        self.mid.block_2 = ResnetBlock(
            block_in,
            block_in,
            dtype=dtype,
            device=device,
        )

        self.up = nn.ModuleList()
        for level in reversed(range(self.num_resolutions)):
            block = nn.ModuleList()
            block_out = config.channels * config.channel_multipliers[level]
            for _ in range(self.num_res_blocks + 1):
                block.append(
                    ResnetBlock(
                        block_in,
                        block_out,
                        dtype=dtype,
                        device=device,
                    )
                )
                block_in = block_out
            up = nn.Module()
            up.block = block
            up.attn = nn.ModuleList()
            if level != 0:
                up.upsample = Upsample(block_in, dtype=dtype, device=device)
            self.up.insert(0, up)

        self.norm_out = _normalize(block_in).to(device=device, dtype=dtype)
        self.conv_out = nn.Conv2d(
            block_in,
            config.out_channels,
            kernel_size=3,
            padding=1,
            dtype=dtype,
            device=device,
        )

    def forward(self, latent: torch.Tensor) -> torch.Tensor:
        """Decode a VAE latent tensor into an image tensor in [-1, 1]."""

        h = self.conv_in(latent)
        h = self.mid.block_1(h)
        h = self.mid.attn_1(h)
        h = self.mid.block_2(h)

        for level in reversed(range(self.num_resolutions)):
            for block in self.up[level].block:
                h = block(h)
            if level != 0:
                h = self.up[level].upsample(h)

        h = self.norm_out(h)
        h = F.silu(h)
        return self.conv_out(h)


class ZImageVae:
    """Load the local Z-Image VAE and decode final latents."""

    def __init__(
        self,
        model_path: ModelSource | Path | None,
        device,
        dtype,
        *,
        model: nn.Module | None = None,
        decoder_config: VaeConfig | None = None,
    ) -> None:
        self.device = device
        self.dtype = dtype
        if model is not None:
            self.model = model.to(device=device, dtype=dtype).eval()
            return
        if model_path is None:
            raise ValueError("model_path is required when model is absent")

        if isinstance(model_path, ModelSource):
            model_source = model_path
        else:
            model_source = ModelSource(model_path, VAE_ROLE)

        state_dict = loadStateDict(model_source)
        config = decoder_config or detectVaeConfig(state_dict)
        decoder = VaeDecoder(config, dtype=dtype, device=device)
        decoder_state = {
            key.removeprefix("decoder."): value
            for key, value in state_dict.items()
            if key.startswith("decoder.")
        }
        missing, unexpected = decoder.load_state_dict(
            decoder_state,
            strict=False,
        )
        if missing or unexpected:
            raise DiffusionCliError(
                "Unsupported VAE state dict: decoder keys did not match"
            )
        self.model = decoder.eval()

    @torch.inference_mode()
    def decode(self, latent: torch.Tensor) -> torch.Tensor:
        """Decode sampled latents into NCHW image tensors in [0, 1]."""

        if latent.ndim != 4:
            raise DiffusionCliError(
                f"Expected NCHW latent tensor, got {tuple(latent.shape)}"
            )
        if latent.shape[1] != LATENT_CHANNELS:
            raise DiffusionCliError(
                f"Expected {LATENT_CHANNELS} latent channels, "
                f"got {latent.shape[1]}"
            )

        latent = latent.to(device=self.device, dtype=self.dtype)
        image = self.model(fluxLatentToVae(latent))
        expected_shape = (
            latent.shape[0],
            VAE_OUT_CHANNELS,
            latent.shape[2] * LATENT_DOWNSCALE,
            latent.shape[3] * LATENT_DOWNSCALE,
        )
        if tuple(image.shape) != expected_shape:
            raise DiffusionCliError(
                "VAE decode produced unexpected shape: "
                f"expected {expected_shape}, got {tuple(image.shape)}"
            )
        if not torch.all(torch.isfinite(image)):
            raise DiffusionCliError("VAE decode produced NaN or Inf values")
        return postprocessVaeOutput(image)

    def toDevice(self, device, dtype) -> None:
        """Move the VAE decoder to the active generation device."""

        self.device = device
        self.dtype = dtype
        self.model.to(device=device, dtype=dtype)
        self.model.eval()

    def toCpu(self) -> None:
        """Move the VAE decoder back to CPU memory."""

        self.device = torch.device("cpu")
        self.model.to(device=self.device)
        self.model.eval()