"""Direct prompt-only MiniMax H3 T2V shape and latent helpers."""

import torch


FPS = 24
AUDIO_LATENT_FPS = 40


def align_frame_count(frames: int) -> int:
    """Snap to H3's valid 17k+5 temporal grid."""
    frames = max(5, frames)
    return frames + (5 - frames) % 17


def temporal_shape(frames: int) -> tuple[int, int, int]:
    """Return output frames, video-latent frames, and joint audio-latent steps."""
    frames = align_frame_count(frames)
    video_steps = 2 if frames <= 5 else ((frames - 5) // 17) * 5 + 2
    return frames, video_steps, round(frames / FPS * AUDIO_LATENT_FPS)


def empty_av_latents(width: int, height: int, frames: int, *, device: torch.device | str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
    """Allocate H3's joint video/audio sampling state without ComfyUI objects."""
    if width % 32 or height % 32:
        raise ValueError("H3 T2V dimensions must be multiples of 32.")
    frames, video_steps, audio_steps = temporal_shape(frames)
    video = torch.zeros((1, 24, video_steps, height // 16, width // 16), device=device)
    audio = torch.zeros((1, 32, 2, audio_steps), device=device)
    return video, audio, frames


def random_av_latents(width: int, height: int, frames: int, seed: int, *, device: torch.device | str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
    """Generate Comfy-equivalent CPU-seeded joint AV noise latents."""
    video, audio, frames = empty_av_latents(width, height, frames, device="cpu")
    generator = torch.manual_seed(seed)
    video = torch.randn(video.shape, dtype=torch.float32, generator=generator, device="cpu").to(video.dtype)
    audio = torch.randn(audio.shape, dtype=torch.float32, generator=generator, device="cpu").to(audio.dtype)
    return video.to(device), audio.to(device), frames