import folder_paths
from comfy import model_management as mm
from comfy.utils import load_torch_file, common_upscale
from accelerate import init_empty_weights
import torch
from ..utils import log


class MultiTalkModelLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (folder_paths.get_filename_list("unet_gguf") + folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}),
            },
        }

    RETURN_TYPES = ("MULTITALKMODEL",)
    RETURN_NAMES = ("model", )
    FUNCTION = "loadmodel"
    CATEGORY = "WanVideoWrapper"

    def loadmodel(self, model, base_precision=None):
        from .multitalk import AudioProjModel
        offload_device = mm.unet_offload_device()
        
        model_path = folder_paths.get_full_path_or_raise("diffusion_models", model)
        if model_path.endswith(".gguf"):
            from diffusers.models.model_loading_utils import load_gguf_checkpoint
            sd = load_gguf_checkpoint(model_path)
        else:
            sd = load_torch_file(model_path, device=offload_device, safe_load=True)

        audio_window=5
        intermediate_dim=512
        output_dim=768
        context_tokens=32
        vae_scale=4
        norm_output_audio = True

        with init_empty_weights():
            multitalk_proj_model = AudioProjModel(
                    seq_len=audio_window,
                    seq_len_vf=audio_window+vae_scale-1,
                    intermediate_dim=intermediate_dim,
                    output_dim=output_dim,
                    context_tokens=context_tokens,
                    norm_output_audio=norm_output_audio,
            )

        multitalk = {
            "proj_model": multitalk_proj_model,
            "sd": sd,
            "is_gguf": model_path.endswith(".gguf"),
            "model_type": "InfiniteTalk" if "infinite" in model.lower() else "MultiTalk",
        }

        return (multitalk,)

def loudness_norm(audio_array, sr=16000, lufs=-23):
    try:
        import pyloudnorm
    except:
        raise ImportError("pyloudnorm package is not installed")
    meter = pyloudnorm.Meter(sr)
    loudness = meter.integrated_loudness(audio_array)
    if abs(loudness) > 100:
        return audio_array
    normalized_audio = pyloudnorm.normalize.loudness(audio_array, loudness, lufs)
    return normalized_audio
    
class MultiTalkWav2VecEmbeds:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "wav2vec_model": ("WAV2VECMODEL",),
            "audio_1": ("AUDIO",),
            "normalize_loudness": ("BOOLEAN", {"default": True, "tooltip": "Normalize the audio loudness to -23 LUFS"}),
            "num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 1, "tooltip": "The total frame count to generate."}),
            "fps": ("FLOAT", {"default": 25.0, "min": 1.0, "max": 60.0, "step": 0.1}),
            "audio_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the audio conditioning"}),
            "audio_cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "When not 1.0, an extra model pass without audio conditioning is done: slower inference but more motion is allowed"}),
            "multi_audio_type": (["para", "add"], {"default": "para", "tooltip": "'para' overlay speakers in parallel, 'add' concatenate sequentially"}),
        },
            "optional" : {
                "audio_2": ("AUDIO",),
                "audio_3": ("AUDIO",),
                "audio_4": ("AUDIO",),
                "ref_target_masks": ("MASK", {"tooltip": "Per-speaker semantic mask(s) in pixel space. Supply one mask per speaker (plus optional background) to guide mouth assignment"}),
            }
        }

    RETURN_TYPES = ("MULTITALK_EMBEDS", "AUDIO", "INT", )
    RETURN_NAMES = ("multitalk_embeds", "audio", "num_frames", )
    FUNCTION = "process"
    CATEGORY = "WanVideoWrapper"

    def process(self, wav2vec_model, normalize_loudness, fps, num_frames, audio_1, audio_scale, audio_cfg_scale, multi_audio_type, audio_2=None, audio_3=None, audio_4=None, ref_target_masks=None):
        model_type = wav2vec_model["model_type"]
        if not "tencent" in model_type.lower():
            raise ValueError("Only tencent wav2vec models supported by MultiTalk")
        import torchaudio
        import numpy as np
        from einops import rearrange

        device = mm.get_torch_device()
        offload_device = mm.unet_offload_device()
        dtype = wav2vec_model["dtype"]
        wav2vec = wav2vec_model["model"]
        wav2vec_feature_extractor = wav2vec_model["feature_extractor"]

        sr = 16000

        audio_inputs = [audio_1, audio_2, audio_3, audio_4]
        audio_inputs = [a for a in audio_inputs if a is not None]

        multitalk_audio_features = []
        seq_lengths = []
        audio_outputs = []  # for debugging / optional saving – choose first as return

        for audio in audio_inputs:
            audio_input = audio["waveform"]
            sample_rate = audio["sample_rate"]

            if sample_rate != 16000:
                audio_input = torchaudio.functional.resample(audio_input, sample_rate, sr)
            audio_input = audio_input[0][0]

            start_time = 0
            end_time = num_frames / fps

            start_sample = int(start_time * sr)
            end_sample = int(end_time * sr)

            try:
                audio_segment = audio_input[start_sample:end_sample]
            except Exception:
                audio_segment = audio_input

            audio_segment = audio_segment.numpy()

            if normalize_loudness:
                audio_segment = loudness_norm(audio_segment, sr=sr)

            audio_feature = np.squeeze(
                wav2vec_feature_extractor(audio_segment, sampling_rate=sr).input_values
            )

            audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
            audio_feature = audio_feature.unsqueeze(0)

            # audio encoder
            audio_duration = len(audio_segment) / sr
            video_length = audio_duration * fps

            wav2vec.to(device)
            embeddings = wav2vec(audio_feature.to(dtype), seq_len=int(video_length), output_hidden_states=True)
            wav2vec.to(offload_device)

            if len(embeddings) == 0:
                print("Fail to extract audio embedding for one speaker")
                continue

            audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
            audio_emb = rearrange(audio_emb, "b s d -> s b d")

            multitalk_audio_features.append(audio_emb.cpu().detach())
            seq_lengths.append(audio_emb.shape[0])

            waveform_tensor = torch.from_numpy(audio_segment).float().cpu().unsqueeze(0).unsqueeze(0)  # (B, C, N)
            audio_outputs.append({"waveform": waveform_tensor, "sample_rate": sr})

        log.info("[MultiTalk] --- Raw speaker lengths (samples) ---")
        for idx, ao in enumerate(audio_outputs):
            log.info(f"  speaker {idx+1}: {ao['waveform'].shape[-1]} samples (shape: {ao['waveform'].shape})")

        # Pad / combine depending on multi_audio_type
        if len(multitalk_audio_features) > 1:
            if multi_audio_type == "para":
                max_len = max(seq_lengths)
                padded = []
                for emb in multitalk_audio_features:
                    if emb.shape[0] < max_len:
                        pad = torch.zeros(max_len - emb.shape[0], *emb.shape[1:], dtype=emb.dtype)
                        emb = torch.cat([emb, pad], dim=0)
                    padded.append(emb)
                multitalk_audio_features = padded
            elif multi_audio_type == "add":
                total_len = sum(seq_lengths)
                full_list = []
                offset = 0
                for emb, length in zip(multitalk_audio_features, seq_lengths):
                    full = torch.zeros(total_len, *emb.shape[1:], dtype=emb.dtype)
                    full[offset:offset+length] = emb
                    full_list.append(full)
                    offset += length
                multitalk_audio_features = full_list

        # fallback
        if len(multitalk_audio_features) == 0:
            raise RuntimeError("No valid audio embeddings extracted, please check inputs")

        multitalk_embeds = {
            "audio_features": multitalk_audio_features,
            "audio_scale": audio_scale,
            "audio_cfg_scale": audio_cfg_scale,
            "ref_target_masks": ref_target_masks
        }

        if len(audio_outputs) == 1: # single speaker
            out_audio = audio_outputs[0]
        else: # multi speaker
            if multi_audio_type == "para":
                # Overlay speakers in parallel – mix waveforms to same length (max len)
                max_len = max([a["waveform"].shape[-1] for a in audio_outputs])
                mixed = torch.zeros(1, 1, max_len, dtype=audio_outputs[0]["waveform"].dtype)
                for a in audio_outputs:
                    w = a["waveform"]
                    if w.shape[-1] < max_len:
                        w = torch.nn.functional.pad(w, (0, max_len - w.shape[-1]))
                    mixed += w
                out_audio = {"waveform": mixed, "sample_rate": sr}
            else:  # "add" – sequential concatenate with silent padding for other speakers
                total_len = sum([a["waveform"].shape[-1] for a in audio_outputs])
                mixed = torch.zeros(1, 1, total_len, dtype=audio_outputs[0]["waveform"].dtype)
                offset = 0
                for a in audio_outputs:
                    w = a["waveform"]
                    mixed[:, :, offset:offset + w.shape[-1]] += w
                    offset += w.shape[-1]
                out_audio = {"waveform": mixed, "sample_rate": sr}

        # Calculate actual frames based on audio duration
        actual_num_frames = num_frames
        if len(audio_outputs) > 0:
            if multi_audio_type == "para":
                # For parallel mode, use the longest audio duration
                max_audio_duration = max([ao["waveform"].shape[-1] / sr for ao in audio_outputs])
                actual_frames_from_audio = int(max_audio_duration * fps)
            else:  # "add"
                # For sequential mode, use the total audio duration
                total_audio_duration = sum([ao["waveform"].shape[-1] / sr for ao in audio_outputs])
                actual_frames_from_audio = int(total_audio_duration * fps)
            
            # Use the smaller of requested frames or actual audio frames
            actual_num_frames = min(num_frames, actual_frames_from_audio)
            
            if actual_frames_from_audio < num_frames:
                log.info(f"[MultiTalk] Audio duration ({actual_frames_from_audio} frames) is shorter than requested ({num_frames} frames). Using {actual_num_frames} frames.")

        # Debug: log final mixed audio length and mode
        total_samples_raw = sum([ao["waveform"].shape[-1] for ao in audio_outputs])
        log.info(f"[MultiTalk] total raw duration = {total_samples_raw/sr:.3f}s")
        log.info(f"[MultiTalk] multi_audio_type={multi_audio_type} | final waveform shape={out_audio['waveform'].shape} | length={out_audio['waveform'].shape[-1]} samples | seconds={out_audio['waveform'].shape[-1]/sr:.3f}s (expected {'sum' if multi_audio_type=='add' else 'max'} of raw)")

        return (multitalk_embeds, out_audio, actual_num_frames)


class WanVideoImageToVideoMultiTalk:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "vae": ("WANVAE",),
            "width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the generation"}),
            "height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the generation"}),
            "frame_window_size": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "The number of frames to process at once, should be a value the model is generally good at."}),
            "motion_frame": ("INT", {"default": 25, "min": 1, "max": 10000, "step": 1, "tooltip": "Driven frame length used in the long video generation. Basically the overlap length."}),
            "force_offload": ("BOOLEAN", {"default": False, "tooltip": "Whether to force offload the model within the loop for VAE operations, enable if you encounter memory issues."}),
            "colormatch": (
            [   
                'disabled',
                'mkl',
                'hm', 
                'reinhard', 
                'mvgd', 
                'hm-mvgd-hm', 
                'hm-mkl-hm',
            ], {
               "default": 'disabled', "tooltip": "Color matching method to use between the windows"
            },),
            },
            "optional": {
                "start_image": ("IMAGE", {"tooltip": "Images to encode"}),
                "tiled_vae": ("BOOLEAN", {"default": False, "tooltip": "Use tiled VAE encoding for reduced memory use"}),
                "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
                "mode": ([
                    "auto",
                    "multitalk",
                    "infinitetalk"
                ], {"default": "auto", "tooltip": "The sampling strategy to use in the long video generation loop, should match the model used"})
            }
        }

    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
    RETURN_NAMES = ("image_embeds",)
    FUNCTION = "process"
    CATEGORY = "WanVideoWrapper"
    DESCRIPTION = "Enables Multi/InfiniteTalk long video generation sampling method, the video is created in windows with overlapping frames. Not compatible or necessary to be used with context windows and many other features besides Multi/InfiniteTalk."

    def process(self, vae, width, height, frame_window_size, motion_frame, force_offload, colormatch, start_image=None, tiled_vae=False, clip_embeds=None, mode="multitalk"):

        H = height
        W = width
        VAE_STRIDE = (4, 8, 8)
        
        num_frames = ((frame_window_size - 1) // 4) * 4 + 1

        # Resize and rearrange the input image dimensions
        if start_image is not None:
            resized_start_image = common_upscale(start_image.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(0, 1)
            resized_start_image = resized_start_image * 2 - 1
            resized_start_image = resized_start_image.unsqueeze(0)
        
        target_shape = (16, (num_frames - 1) // VAE_STRIDE[0] + 1,
                        height // VAE_STRIDE[1],
                        width // VAE_STRIDE[2])

        image_embeds = {
            "multitalk_sampling": True,
            "multitalk_start_image": resized_start_image if start_image is not None else None,
            "num_frames": num_frames,
            "motion_frame": motion_frame,
            "target_h": H,
            "target_w": W,
            "tiled_vae": tiled_vae,
            "force_offload": force_offload,
            "vae": vae,
            "target_shape": target_shape,
            "clip_context": clip_embeds.get("clip_embeds", None) if clip_embeds is not None else None,
            "colormatch": colormatch,
            "multitalk_mode": mode
        }

        return (image_embeds,)
    
NODE_CLASS_MAPPINGS = {
    "MultiTalkModelLoader": MultiTalkModelLoader,
    "MultiTalkWav2VecEmbeds": MultiTalkWav2VecEmbeds,
    "WanVideoImageToVideoMultiTalk": WanVideoImageToVideoMultiTalk    
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MultiTalkModelLoader": "Multi/InfiniteTalk Model Loader",
    "MultiTalkWav2VecEmbeds": "Multi/InfiniteTalk Wav2Vec Embeds",
    "WanVideoImageToVideoMultiTalk": "WanVideo Long I2V Multi/InfiniteTalk"
}