import torch

import node_helpers
from comfy_extras.nodes_wan import get_audio_embed_bucket_fps, linear_interpolation


def apply_wan_s2v_audio_conditioning(positive, negative, length, audio_encoder_output=None, frame_offset=0):
    if audio_encoder_output is None:
        return positive, negative, frame_offset

    latent_t = ((length - 1) // 4) + 1
    feat = torch.cat(audio_encoder_output["encoded_audio_all_layers"])
    video_rate = 30
    fps = 16
    feat = linear_interpolation(feat, input_fps=50, output_fps=video_rate)
    batch_frames = latent_t * 4
    audio_embed_bucket, _ = get_audio_embed_bucket_fps(feat, fps=fps, batch_frames=batch_frames, m=0, video_rate=video_rate)
    audio_embed_bucket = audio_embed_bucket.unsqueeze(0)
    if len(audio_embed_bucket.shape) == 3:
        audio_embed_bucket = audio_embed_bucket.permute(0, 2, 1)
    elif len(audio_embed_bucket.shape) == 4:
        audio_embed_bucket = audio_embed_bucket.permute(0, 2, 3, 1)

    audio_embed_bucket = audio_embed_bucket[:, :, :, frame_offset:frame_offset + batch_frames]
    if audio_embed_bucket.shape[3] > 0:
        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})
        frame_offset += batch_frames
    return positive, negative, frame_offset