#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
from typing import Any

import torch


FORMAT_VERSION = 1


def prepare_inference_model(model: torch.nn.Module) -> torch.nn.Module:
    return model.eval().requires_grad_(False)


def save_latent_bundle(
    path: Path,
    *,
    video_latents: torch.Tensor,
    audio_latents: torch.Tensor,
    metadata: dict[str, Any],
) -> None:
    if video_latents.ndim != 5 or video_latents.shape[0] != 1:
        raise ValueError(
            "video latents must have shape (1, channels, frames, height, width)"
        )
    if audio_latents.ndim != 3 or audio_latents.shape[0] != 2:
        raise ValueError("audio latents must have shape (2, channels, samples)")

    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    torch.save(
        {
            "format_version": FORMAT_VERSION,
            "metadata": dict(metadata),
            "video_latents": video_latents.detach().to("cpu").contiguous(),
            "audio_latents": audio_latents.detach().to("cpu").contiguous(),
        },
        temporary,
    )
    temporary.replace(path)


def load_latent_bundle(path: Path) -> dict[str, Any]:
    payload = torch.load(path, map_location="cpu", weights_only=False)
    if payload.get("format_version") != FORMAT_VERSION:
        raise ValueError(f"unsupported latent bundle format: {payload.get('format_version')}")
    return payload