import torch

import comfy.utils

WAN_VAE_SCALE = 8
WAN_PATCH_SPATIAL = 2


def _padded_latent_dim(pixels: int) -> int:
    latent = pixels // WAN_VAE_SCALE
    return latent + (WAN_PATCH_SPATIAL - latent % WAN_PATCH_SPATIAL) % WAN_PATCH_SPATIAL


def token_grid_size(width: int, height: int) -> tuple[int, int]:
    return _padded_latent_dim(height) // WAN_PATCH_SPATIAL, _padded_latent_dim(width) // WAN_PATCH_SPATIAL


def mask_to_token_grid(mask_image: torch.Tensor, width: int, height: int) -> torch.Tensor:
    token_h, token_w = token_grid_size(width, height)
    mask = mask_image[0] if mask_image.ndim == 3 else mask_image
    mask = mask.unsqueeze(0).unsqueeze(0)
    mask = comfy.utils.common_upscale(mask, width, height, "area", "center")
    mask = comfy.utils.common_upscale(mask, token_w, token_h, "area", "center")
    return (mask > 0.5).to(dtype=torch.float32).flatten(2).squeeze(1)


def _latent_frame_weight(video_frame: int, start_frame: int, end_frame: int, crossfade_frames: int) -> float:
    if video_frame < start_frame or video_frame >= end_frame:
        return 0.0
    if crossfade_frames <= 0:
        return 1.0
    weight = 1.0
    if video_frame < start_frame + crossfade_frames:
        weight = min(weight, (video_frame - start_frame + 1) / crossfade_frames)
    if video_frame >= end_frame - crossfade_frames:
        weight = min(weight, (end_frame - video_frame) / crossfade_frames)
    return max(0.0, weight)


def build_timeline_audio_inject_mask(
    width: int,
    height: int,
    length: int,
    segments,
    crossfade_frames: int = 0,
    device=None,
) -> torch.Tensor:
    latent_t = ((length - 1) // 4) + 1
    token_h, token_w = token_grid_size(width, height)
    n_tokens = token_h * token_w
    mask = torch.zeros(1, latent_t, n_tokens, 1)

    for segment in segments:
        tokens = mask_to_token_grid(segment["mask_image"], width, height)
        start_frame = int(segment["start_frame"])
        end_frame = int(segment["end_frame"])
        for latent_idx in range(latent_t):
            vf0 = latent_idx * 4
            vf1 = vf0 + 4
            weight = 0.0
            for video_frame in range(vf0, vf1):
                weight = max(weight, _latent_frame_weight(video_frame, start_frame, end_frame, crossfade_frames))
            if weight > 0.0:
                mask[:, latent_idx, :, 0] = torch.maximum(mask[:, latent_idx, :, 0], tokens * weight)

    if device is not None:
        mask = mask.to(device)
    return mask