#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Author: C0ffymachyne
License: GPLv3
Version: 1.0.0

Description:
    Mixdown node for pad synths
"""

import torch

from typing import Tuple, List, Dict, Union, Any
import torchaudio

from ..core.io import audio_to_comfy_3d
from ..core.loudness import lufs_normalization


class SignalProcessingMixdown:
    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Any]:
        return {
            "required": {
                "audio_inputs": ("AUDIO_LIST", {"default": []}),
            },
            "optional": {
                "gain_factors": (
                    "FLOAT_LIST",
                    {"default": [], "min": 0.0, "max": 2.0, "step": 0.1},
                ),
                # If empty, default to [1.0] * num_audios
            },
        }

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("mixed_audio",)
    CATEGORY = "Signal Processing"
    FUNCTION = "process"

    def process(
        self,
        audio_inputs: List[Dict[str, Union[torch.Tensor, int]]],
        gain_factors: List[float] = [],
    ) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
        """
        Mix down multiple audio inputs into a single audio output
        with optional individual volume controls.

        Parameters:
            audio_inputs (List[Dict]): List of audio inputs,
            each containing 'waveform' and 'sample_rate'.
            output_normalization (float): Normalization factor for the mixed audio (0.0 to 1.0).
            gain_factors (List[float], optional): List of gain factors for each audio input.

        Returns:
            Tuple[Dict[str, torch.Tensor], int]: Mixed audio with waveform and sample rate.
        """

        if not audio_inputs:
            raise ValueError("No audio inputs provided for mixing.")

        num_audios = len(audio_inputs)

        # Handle gain_factors
        if not gain_factors:
            gain_factors = [1.0] * num_audios
        elif len(gain_factors) != num_audios:
            raise ValueError(
                f"Number of gain factors ({len(gain_factors)}) \
                    does not match number of audio inputs ({num_audios})."
            )

        # Extract sample rates and verify consistency
        sample_rates: List[int] = [audio["sample_rate"] for audio in audio_inputs]
        target_sample_rate = sample_rates[0]

        for idx, sr in enumerate(sample_rates):
            if sr != target_sample_rate:
                resampler = torchaudio.transforms.Resample(
                    orig_freq=sr, new_freq=target_sample_rate
                )

                _waveform: torch.Tensor = audio_inputs[idx]["waveform"]
                resampler.to(
                    device=_waveform.device,
                    dtype=_waveform.dtype,
                )
                audio_inputs[idx]["waveform"] = resampler(_waveform)
                audio_inputs[idx]["sample_rate"] = target_sample_rate

        # Determine the maximum length among all audio inputs
        lengthsw: List[torch.Tensor] = [audio["waveform"] for audio in audio_inputs]
        lengths: List[int] = [wave.shape[-1] for wave in lengthsw]
        max_length = max(lengths)

        # Pad or truncate each audio to match the maximum length and apply gain
        for idx, audio in enumerate(audio_inputs):
            waveform: torch.Tensor = audio["waveform"]
            current_length = waveform.shape[-1]
            gain = gain_factors[idx]

            if current_length < max_length:
                padding = max_length - current_length
                # Pad with zeros (silence) at the end
                waveform = torch.nn.functional.pad(waveform, (0, padding))
            elif current_length > max_length:
                # Truncate the waveform to max_length
                waveform = waveform[:, :, :max_length]

            # Apply gain
            waveform = waveform * gain

            audio["waveform"] = waveform

        # Sum all waveforms to create the mix
        mixed_waveform: torch.Tensor = torch.zeros_like(audio_inputs[0]["waveform"])
        for idx, audio in enumerate(audio_inputs):
            mixed_waveform += audio["waveform"]

        mixed_waveform = lufs_normalization(mixed_waveform, target_sample_rate)

        return audio_to_comfy_3d(mixed_waveform, target_sample_rate)