"""NormalCrafter inference pipeline.

Full-video intermediates stay on CPU. Only the active CLIP/VAE chunk or UNet
window is moved to the inference device.
"""

from __future__ import annotations

from dataclasses import dataclass
import inspect
import math
from typing import Callable

import torch
import torch.nn.functional as F

from .model import NormalCrafterModels

FPS_CONDITIONING = 7
MOTION_BUCKET_ID = 127
NOISE_AUG_STRENGTH = 0.0

ProgressCallback = Callable[[], None]


@dataclass
class FrameSource:
    """Lazy CPU-side preprocessing view over a ComfyUI IMAGE batch."""

    frames: torch.Tensor
    max_resolution: int
    window_size: int

    def __post_init__(self) -> None:
        if self.frames.ndim != 4 or self.frames.shape[-1] < 3:
            raise ValueError("images must have shape [frames, height, width, channels>=3]")
        if self.frames.shape[0] == 0:
            raise ValueError("images contains no frames")

        self.frames = (
            self.frames[..., :3]
            .detach()
            .to(device="cpu", dtype=torch.float32)
            .contiguous()
            .clamp_(0.0, 1.0)
        )

        _, self.original_height, self.original_width, _ = self.frames.shape
        scale = min(
            1.0,
            self.max_resolution / max(self.original_height, self.original_width),
        )
        self.height = max(1, int(round(self.original_height * scale)))
        self.width = max(1, int(round(self.original_width * scale)))

        padded_height = math.ceil(self.height / 64) * 64
        padded_width = math.ceil(self.width / 64) * 64
        pad_h = padded_height - self.height
        pad_w = padded_width - self.width
        self.pad_top = pad_h // 2
        self.pad_bottom = pad_h - self.pad_top
        self.pad_left = pad_w // 2
        self.pad_right = pad_w - self.pad_left

    @property
    def frame_count(self) -> int:
        return int(self.frames.shape[0])

    @property
    def effective_frame_count(self) -> int:
        return max(self.frame_count, self.window_size)

    def get(self, start: int, end: int) -> torch.Tensor:
        """Return processed BCHW frames; short clips repeat their final frame."""

        indices = torch.arange(start, end, dtype=torch.long).clamp_max(
            self.frame_count - 1
        )
        chunk = self.frames.index_select(0, indices).permute(0, 3, 1, 2).contiguous()

        if (self.height, self.width) != (self.original_height, self.original_width):
            chunk = F.interpolate(
                chunk,
                size=(self.height, self.width),
                mode="bicubic",
                align_corners=False,
                antialias=True,
            ).clamp_(0.0, 1.0)

        if self.pad_top or self.pad_bottom or self.pad_left or self.pad_right:
            chunk = F.pad(
                chunk,
                (self.pad_left, self.pad_right, self.pad_top, self.pad_bottom),
                value=1.0,
            )
        return chunk

    def crop(self, frames: torch.Tensor) -> torch.Tensor:
        return frames[
            :,
            :,
            self.pad_top : self.pad_top + self.height,
            self.pad_left : self.pad_left + self.width,
        ]


def _validate_settings(window_size: int, step_size: int, chunk_size: int) -> None:
    if window_size < 2:
        raise ValueError("window_size must be at least 2")
    if not 1 <= step_size <= window_size:
        raise ValueError("step_size must be between 1 and window_size")
    if chunk_size < 1:
        raise ValueError("chunk_size must be at least 1")


def _build_windows(
    frame_count: int,
    window_size: int,
    step_size: int,
) -> list[tuple[int, int]]:
    """Build the same full-size, tail-anchored windows as NormalCrafter."""

    if frame_count < window_size:
        raise ValueError("frame_count must already be padded to window_size")

    last_start = frame_count - window_size
    starts = list(range(0, last_start + 1, step_size))
    if starts[-1] != last_start:
        starts.append(last_start)
    return [(start, start + window_size) for start in starts]


def estimate_progress_steps(
    frame_count: int,
    *,
    window_size: int,
    step_size: int,
    chunk_size: int,
) -> int:
    _validate_settings(window_size, step_size, chunk_size)
    effective = max(frame_count, window_size)
    return max(
        1,
        math.ceil(effective / chunk_size)
        + math.ceil(effective / chunk_size)
        + len(_build_windows(effective, window_size, step_size))
        + math.ceil(frame_count / chunk_size),
    )


def run_inference(
    models: NormalCrafterModels,
    images: torch.Tensor,
    *,
    device: torch.device,
    max_resolution: int,
    window_size: int,
    step_size: int,
    chunk_size: int,
    progress: ProgressCallback | None = None,
) -> torch.Tensor:
    """Run CLIP -> VAE encode -> temporal UNet -> VAE decode."""

    _validate_settings(window_size, step_size, chunk_size)
    source = FrameSource(
        images,
        max_resolution=max_resolution,
        window_size=window_size,
    )

    try:
        with torch.inference_mode():
            embeddings = _encode_clip(models, source, device, chunk_size, progress)
            rgb_latents = _encode_rgb(models, source, device, chunk_size, progress)
            normal_latents = _predict_normals(
                models,
                rgb_latents,
                embeddings,
                device,
                window_size,
                step_size,
                progress,
            )
            del rgb_latents, embeddings
            return _decode_normals(
                models,
                normal_latents,
                source,
                device,
                chunk_size,
                progress,
            )
    finally:
        models.offload_all()


def _tick(progress: ProgressCallback | None) -> None:
    if progress is not None:
        progress()


def _encode_clip(
    models: NormalCrafterModels,
    source: FrameSource,
    device: torch.device,
    chunk_size: int,
    progress: ProgressCallback | None,
) -> torch.Tensor:
    models.activate("image_encoder", device)

    mean = torch.tensor(
        models.feature_extractor.image_mean,
        dtype=torch.float32,
    ).view(1, 3, 1, 1)
    std = torch.tensor(
        models.feature_extractor.image_std,
        dtype=torch.float32,
    ).view(1, 3, 1, 1)
    output = None

    for start in range(0, source.effective_frame_count, chunk_size):
        end = min(start + chunk_size, source.effective_frame_count)
        frames = source.get(start, end)
        frames = F.interpolate(
            frames,
            size=(224, 224),
            mode="bicubic",
            align_corners=False,
            antialias=True,
        ).clamp_(0.0, 1.0)

        pixels = ((frames - mean) / std).to(device=device, dtype=models.dtype)
        embeddings = models.image_encoder(pixel_values=pixels).image_embeds
        if embeddings.ndim == 2:
            embeddings = embeddings.unsqueeze(1)
        embeddings = embeddings.detach().to(device="cpu", dtype=models.dtype)

        if output is None:
            output = torch.empty(
                (source.effective_frame_count, *embeddings.shape[1:]),
                dtype=models.dtype,
                device="cpu",
            )
        output[start:end].copy_(embeddings)
        _tick(progress)

    assert output is not None
    return output


def _encode_rgb(
    models: NormalCrafterModels,
    source: FrameSource,
    device: torch.device,
    chunk_size: int,
    progress: ProgressCallback | None,
) -> torch.Tensor:
    models.activate("vae", device)
    vae_dtype = next(models.vae.parameters()).dtype
    output = None

    for start in range(0, source.effective_frame_count, chunk_size):
        end = min(start + chunk_size, source.effective_frame_count)
        frames = (source.get(start, end) * 2.0 - 1.0).to(
            device=device,
            dtype=vae_dtype,
        )
        latents = models.vae.encode(frames).latent_dist.mode()
        latents = latents.detach().to(device="cpu", dtype=models.dtype)

        if output is None:
            output = torch.empty(
                (source.effective_frame_count, *latents.shape[1:]),
                dtype=models.dtype,
                device="cpu",
            )
        output[start:end].copy_(latents)
        _tick(progress)

    assert output is not None
    return output


def _added_time_ids(
    models: NormalCrafterModels,
    device: torch.device,
) -> torch.Tensor:
    values = [FPS_CONDITIONING, MOTION_BUCKET_ID, NOISE_AUG_STRENGTH]
    expected = models.unet.add_embedding.linear_1.in_features
    actual = int(models.unet.config.addition_time_embed_dim) * len(values)
    if expected != actual:
        raise ValueError(
            f"UNet added-time embedding mismatch: expected {expected}, got {actual}"
        )
    return torch.tensor([values], dtype=models.dtype, device=device)


def _predict_normals(
    models: NormalCrafterModels,
    rgb_latents: torch.Tensor,
    embeddings: torch.Tensor,
    device: torch.device,
    window_size: int,
    step_size: int,
    progress: ProgressCallback | None,
) -> torch.Tensor:
    models.activate("unet", device)
    windows = _build_windows(
        int(rgb_latents.shape[0]),
        window_size,
        step_size,
    )
    output = torch.empty_like(rgb_latents, device="cpu")

    previous_window: tuple[int, int] | None = None
    previous_prediction: torch.Tensor | None = None

    for start, end in windows:
        image_latents = rgb_latents[start:end].to(
            device=device,
            dtype=models.dtype,
        ).unsqueeze(0)
        frame_embeddings = embeddings[start:end].to(
            device=device,
            dtype=models.dtype,
        )
        latents = torch.zeros(
            (
                1,
                end - start,
                int(models.unet.config.out_channels),
                *image_latents.shape[-2:],
            ),
            device=device,
            dtype=models.dtype,
        )

        overlap = 0 if previous_window is None else max(0, previous_window[1] - start)
        old_overlap = None
        if overlap and previous_window is not None and previous_prediction is not None:
            previous_offset = start - previous_window[0]
            old_overlap = previous_prediction[
                :,
                previous_offset : previous_offset + overlap,
            ]
            latents[:, :overlap].copy_(old_overlap)

        models.scheduler.set_timesteps(1, device=device)
        timestep = models.scheduler.timesteps[0]
        model_input = models.scheduler.scale_model_input(latents, timestep)
        model_input = torch.cat((model_input, image_latents), dim=2)

        noise_prediction = models.unet(
            model_input,
            timestep,
            encoder_hidden_states=frame_embeddings,
            added_time_ids=_added_time_ids(models, device),
            return_dict=False,
        )[0]
        predicted = models.scheduler.step(
            noise_prediction,
            timestep,
            latents,
        ).prev_sample

        if old_overlap is not None:
            weights = torch.linspace(
                1.0,
                0.0,
                overlap + 2,
                device=device,
                dtype=models.dtype,
            )[1:-1].view(1, overlap, 1, 1, 1)
            predicted[:, :overlap] = (
                old_overlap * weights
                + predicted[:, :overlap] * (1.0 - weights)
            )

        output[start:end].copy_(predicted[0].detach().to("cpu"))
        previous_window = (start, end)
        previous_prediction = predicted.detach()
        _tick(progress)

    return output


def _decode_normals(
    models: NormalCrafterModels,
    normal_latents: torch.Tensor,
    source: FrameSource,
    device: torch.device,
    chunk_size: int,
    progress: ProgressCallback | None,
) -> torch.Tensor:
    models.activate("vae", device)
    vae_dtype = next(models.vae.parameters()).dtype

    output = torch.empty(
        (source.frame_count, source.original_height, source.original_width, 3),
        dtype=torch.float32,
        device="cpu",
    )
    scaling_factor = float(models.vae.config.scaling_factor)

    forward = (
        models.vae._orig_mod.forward
        if hasattr(models.vae, "_orig_mod")
        else models.vae.forward
    )
    accepts_num_frames = "num_frames" in inspect.signature(forward).parameters

    for start in range(0, source.frame_count, chunk_size):
        end = min(start + chunk_size, source.frame_count)
        latents = normal_latents[start:end].to(
            device=device,
            dtype=vae_dtype,
        ) / scaling_factor
        kwargs = {"num_frames": end - start} if accepts_num_frames else {}
        normals = models.vae.decode(latents, **kwargs).sample.float().to("cpu")
        normals = source.crop(normals)

        if (source.height, source.width) != (
            source.original_height,
            source.original_width,
        ):
            normals = F.interpolate(
                normals,
                size=(source.original_height, source.original_width),
                mode="bilinear",
                align_corners=False,
            )

        normals = normals.clamp(-1.0, 1.0)
        normals = F.normalize(normals, p=2.0, dim=1, eps=1e-6)
        normals = (normals * 0.5 + 0.5).clamp_(0.0, 1.0)
        output[start:end].copy_(
            normals.permute(0, 2, 3, 1).contiguous()
        )
        _tick(progress)

    return output