# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
from functools import cache, lru_cache

import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
import torchaudio.compliance.kaldi as kaldi

from vllm_omni.utils.audio import mel_filter_bank

logger = logging.getLogger(__name__)


def dynamic_range_compression_torch(x, c=1, clip_val=1e-5):
    return torch.log(torch.clamp(x, min=clip_val) * c)


def spectral_normalize_torch(magnitudes):
    output = dynamic_range_compression_torch(magnitudes)
    return output


@lru_cache
def _get_mel_basis(
    sampling_rate: int,
    n_fft: int,
    num_mels: int,
    fmin: float,
    fmax: float | None,
    device_str: str,
) -> torch.Tensor:
    return mel_filter_bank(
        sr=sampling_rate,
        n_fft=n_fft,
        n_mels=num_mels,
        fmin=fmin,
        fmax=fmax,
    ).to(torch.device(device_str))


@lru_cache
def _get_hann_window(win_size: int, device_str: str) -> torch.Tensor:
    return torch.hann_window(win_size).to(torch.device(device_str))


def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
    device_str = str(y.device)
    mel = _get_mel_basis(
        int(sampling_rate),
        int(n_fft),
        int(num_mels),
        float(fmin),
        fmax,
        device_str,
    )
    window = _get_hann_window(int(win_size), device_str)

    y = torch.nn.functional.pad(
        y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect"
    )
    y = y.squeeze(1)

    spec = torch.view_as_real(
        torch.stft(
            y,
            n_fft,
            hop_length=hop_size,
            win_length=win_size,
            window=window,
            center=center,
            pad_mode="reflect",
            normalized=False,
            onesided=True,
            return_complex=True,
        )
    )

    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))

    spec = torch.matmul(mel, spec)
    spec = spectral_normalize_torch(spec)

    return spec


def load_wav(wav, target_sr, min_sr=16000):
    if not isinstance(wav, tuple):
        speech, sample_rate = torchaudio.load(wav, backend="soundfile")
    else:
        speech, sample_rate = wav
        if isinstance(speech, np.ndarray):
            speech = torch.from_numpy(speech).float().unsqueeze(0)

    if sample_rate != target_sr:
        if sample_rate < min_sr:
            raise ValueError(
                f"Audio sample rate {sample_rate} Hz is too low. "
                f"Minimum required: {min_sr} Hz, target: {target_sr} Hz. "
                f"Please provide audio with sample rate >= {min_sr} Hz."
            )
        speech = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sr)(speech)

    speech = speech.to(dtype=torch.float32)
    return speech


def extract_speech_feat(prompt_wav, feat_extractor, device):
    speech = load_wav(prompt_wav, 24000)
    speech_feat = feat_extractor(speech).squeeze(dim=0).transpose(0, 1).to(device)
    speech_feat = speech_feat.unsqueeze(dim=0)
    speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32).to(device)
    return speech_feat, speech_feat_len


# Adopted from https://github.com/openai/whisper/blob/main/whisper/audio.py


def exact_div(x, y):
    assert x % y == 0
    return x // y


@cache
def mel_filters(device, n_mels: int) -> torch.Tensor:
    """Compute mel filterbank matrix for projecting STFT into a Mel spectrogram."""
    return mel_filter_bank(sr=16000, n_fft=400, n_mels=n_mels).to(device)


def log_mel_spectrogram(
    audio: str | np.ndarray | torch.Tensor,
    n_mels: int = 80,
    padding: int = 0,
    device: str | torch.device | None = None,
):
    """
    Compute the log-Mel spectrogram of

    Parameters
    ----------
    audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
        The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz

    n_mels: int
        The number of Mel-frequency filters, only 80 and 128 are supported

    padding: int
        Number of zero samples to pad to the right

    device: Optional[Union[str, torch.device]]
        If given, the audio tensor is moved to this device before STFT

    Returns
    -------
    torch.Tensor, shape = (n_mels, n_frames)
        A Tensor that contains the Mel spectrogram
    """
    N_FFT = 400
    HOP_LENGTH = 160

    if not torch.is_tensor(audio):
        raise TypeError(f"audio is not tensor {type(audio)}")

    if device is not None:
        audio = audio.to(device)
    if padding > 0:
        audio = F.pad(audio, (0, padding))
    window = torch.hann_window(N_FFT).to(audio.device)
    stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
    magnitudes = stft[..., :-1].abs() ** 2

    filters = mel_filters(audio.device, n_mels)
    mel_spec = filters @ magnitudes

    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0
    return log_spec


def extract_spk_embedding(prompt_wav, campplus_session, device):
    speech = load_wav(prompt_wav, 16000)
    feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
    feat = feat - feat.mean(dim=0, keepdim=True)
    embedding = (
        campplus_session.run(None, {campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0]
        .flatten()
        .tolist()
    )
    embedding = torch.tensor([embedding]).to(device)
    return embedding


def unpad_prompt_conditioning(speech_token, speech_feat, speech_token_len):
    """Drop right-padding from per-request prompt conditioning.

    The talker emits ``speech_token`` ``[1, T]`` possibly right-padded to the
    batch max, plus the true ``speech_token_len``; ``speech_feat`` is at the 2:1
    mel:token ratio (``[1, 2T, F]``). This trims both to the real length. Called
    in the (eager) code2wav stage where a host read of the length is allowed.
    Returns ``(speech_token, speech_feat)`` unchanged when length is unavailable.
    """
    if not isinstance(speech_token, torch.Tensor) or speech_token_len is None:
        return speech_token, speech_feat
    try:
        token_len = int(torch.as_tensor(speech_token_len).reshape(-1)[0].item())
    except (RuntimeError, ValueError, IndexError):
        return speech_token, speech_feat
    if token_len <= 0 or speech_token.dim() < 2:
        return speech_token, speech_feat
    token_len = min(token_len, int(speech_token.shape[1]))
    speech_token = speech_token[:, :token_len]
    if isinstance(speech_feat, torch.Tensor) and speech_feat.dim() >= 2:
        speech_feat = speech_feat[:, : 2 * token_len]
    return speech_token, speech_feat


def extract_spk_embedding_trt(prompt_wav, campplus_trt, device):
    """TensorRT counterpart of ``extract_spk_embedding``.

    Identical fbank front-end; the campplus forward runs on a prebuilt
    TensorRT engine (GPU) instead of the CPU ONNX-Runtime session. Returns the
    same ``[1, 192]`` embedding tensor on ``device``.
    """
    speech = load_wav(prompt_wav, 16000)
    feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
    feat = feat - feat.mean(dim=0, keepdim=True)
    embedding = campplus_trt(feat)
    return embedding.to(device)


def extract_text_token(text, tokenizer, allowed_special):
    text_token = tokenizer.encode(text, allowed_special=allowed_special)
    text_token = torch.tensor([text_token], dtype=torch.int32)
    text_token_len = text_token.shape[1]
    return text_token, text_token_len


def concat_text_with_prompt_ids(
    text: torch.Tensor,
    text_len: torch.Tensor,
    prompt_text: torch.Tensor,
    prompt_text_len: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    text = torch.concat([prompt_text, text], dim=1)
    text_len = text_len + prompt_text_len
    return text, text_len


def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
    """Make mask tensor containing indices of padded part.

    See description of make_non_pad_mask.

    Args:
        lengths (torch.Tensor): Batch of lengths (B,).
    Returns:
        torch.Tensor: Mask tensor containing indices of padded part.

    Examples:
        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        masks = [[0, 0, 0, 0 ,0],
                 [0, 0, 0, 1, 1],
                 [0, 0, 1, 1, 1]]
    """
    batch_size = lengths.size(0)
    max_len = max_len if max_len > 0 else lengths.max().item()
    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
    seq_length_expand = lengths.unsqueeze(-1)
    mask = seq_range_expand >= seq_length_expand
    return mask