import math

import torch

import node_helpers
from comfy_extras.nodes_wan import get_audio_embed_bucket_fps, linear_interpolation

WAN_AUDIO_INPUT_FPS = 50
WAN_AUDIO_VIDEO_RATE = 30
WAN_AUDIO_FPS = 16
WAN_AUDIO_SAMPLE_RATE = 16000


def _audio_feat(audio_encoder_output):
    feat = torch.cat(audio_encoder_output["encoded_audio_all_layers"])
    return linear_interpolation(feat, input_fps=WAN_AUDIO_INPUT_FPS, output_fps=WAN_AUDIO_VIDEO_RATE)


def audio_encoder_output_video_frames(audio_encoder_output, fps=WAN_AUDIO_FPS):
    audio_samples = audio_encoder_output.get("audio_samples")
    if audio_samples is not None:
        return max(1, int(round(audio_samples / float(WAN_AUDIO_SAMPLE_RATE) * fps)))
    feat = _audio_feat(audio_encoder_output)
    return max(1, int(round(feat.shape[1] * fps / WAN_AUDIO_VIDEO_RATE)))


def _permute_audio_embed_bucket(audio_embed_bucket):
    audio_embed_bucket = audio_embed_bucket.unsqueeze(0)
    if len(audio_embed_bucket.shape) == 3:
        return audio_embed_bucket.permute(0, 2, 1)
    return audio_embed_bucket.permute(0, 2, 3, 1)


def build_timeline_audio_embed(length, segments):
    latent_t = ((length - 1) // 4) + 1
    batch_frames = latent_t * 4
    total_feat_frames = int(math.ceil(batch_frames * WAN_AUDIO_VIDEO_RATE / WAN_AUDIO_FPS))

    composite = None
    cursor_auto = 0
    for segment in segments:
        feat = _audio_feat(segment["audio_encoder_output"])
        if composite is None:
            composite = torch.zeros(
                feat.shape[0], total_feat_frames, feat.shape[2],
                dtype=feat.dtype, device=feat.device)

        start_frame = segment.get("start_frame", -1)
        if start_frame is None or start_frame < 0:
            start_frame = cursor_auto
        else:
            start_frame = int(start_frame)

        start_feat = int(round(start_frame * WAN_AUDIO_VIDEO_RATE / WAN_AUDIO_FPS))
        copy_len = min(feat.shape[1], total_feat_frames - start_feat)
        if copy_len > 0 and start_feat < total_feat_frames:
            composite[:, start_feat:start_feat + copy_len, :] = feat[:, :copy_len, :]

        cursor_auto = start_frame + audio_encoder_output_video_frames(segment["audio_encoder_output"])

    audio_embed_bucket, _ = get_audio_embed_bucket_fps(
        composite, fps=WAN_AUDIO_FPS, batch_frames=batch_frames, m=0, video_rate=WAN_AUDIO_VIDEO_RATE)
    audio_embed_bucket = _permute_audio_embed_bucket(audio_embed_bucket)
    return audio_embed_bucket[:, :, :, :batch_frames]


def resolve_timeline_segment_ranges(length, segments):
    batch_frames = (((length - 1) // 4) + 1) * 4
    resolved = []
    cursor_auto = 0
    for segment in segments:
        start_frame = segment.get("start_frame", -1)
        if start_frame is None or start_frame < 0:
            start_frame = cursor_auto
        else:
            start_frame = int(start_frame)
        end_frame = min(batch_frames, start_frame + audio_encoder_output_video_frames(segment["audio_encoder_output"]))
        resolved.append({**segment, "start_frame": start_frame, "end_frame": end_frame})
        cursor_auto = end_frame
    return resolved


def apply_timeline_audio_conditioning(positive, negative, length, segments):
    audio_embed_bucket = build_timeline_audio_embed(length, segments)
    if audio_embed_bucket is None or audio_embed_bucket.shape[3] <= 0:
        return positive, negative
    positive = node_helpers.conditioning_set_values(positive, {"audio_embed": audio_embed_bucket})
    negative = node_helpers.conditioning_set_values(negative, {"audio_embed": audio_embed_bucket * 0.0})
    return positive, negative