# MIT License
# 
# Copyright (c) 2026 audio-embeddings contributors
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Prepare mono waveforms without changing the learned frontend's numerics."""

from __future__ import annotations

from typing import Any

import numpy as np
import torch
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from transformers.feature_extraction_utils import BatchFeature


class AudioEmbeddingFeatureExtractor(SequenceFeatureExtractor):
    model_input_names = ["input_values", "attention_mask"]

    def __init__(
        self,
        sampling_rate: int = 16000,
        padding_value: float = 0.0,
        return_attention_mask: bool = True,
        **kwargs: Any,
    ) -> None:
        if not isinstance(sampling_rate, int) or sampling_rate <= 0:
            raise ValueError("sampling_rate must be a positive integer")
        kwargs.pop("feature_size", None)
        super().__init__(
            feature_size=1,
            sampling_rate=sampling_rate,
            padding_value=padding_value,
            return_attention_mask=return_attention_mask,
            **kwargs,
        )

    def __call__(
        self,
        raw_speech: Any,
        *,
        sampling_rate: int | None = None,
        padding: bool | str = True,
        max_length: int | None = None,
        truncation: bool = False,
        pad_to_multiple_of: int | None = None,
        return_attention_mask: bool | None = None,
        return_tensors: str | None = "pt",
    ) -> BatchFeature:
        if sampling_rate != self.sampling_rate:
            raise ValueError(
                f"Pass sampling_rate={self.sampling_rate}; got {sampling_rate}. "
                "Resample audio to the model rate before calling the feature extractor."
            )
        if isinstance(raw_speech, torch.Tensor):
            raw_speech = raw_speech.detach().cpu().float().numpy()
        if isinstance(raw_speech, np.ndarray):
            if raw_speech.ndim not in {1, 2}:
                raise ValueError(
                    "Expected mono audio [samples] or a batch [batch, samples]"
                )
            batch = [raw_speech] if raw_speech.ndim == 1 else list(raw_speech)
        elif isinstance(raw_speech, (list, tuple)) and len(raw_speech):
            batch = [raw_speech] if np.isscalar(raw_speech[0]) else list(raw_speech)
        else:
            raise ValueError("Provide a nonempty waveform or batch of mono waveforms")
        waveforms = []
        for waveform in batch:
            if isinstance(waveform, torch.Tensor):
                waveform = waveform.detach().cpu().float().numpy()
            array = np.asarray(waveform, dtype=np.float32)
            if array.ndim != 1 or array.size == 0 or not np.isfinite(array).all():
                raise ValueError(
                    "Each waveform must be a nonempty, finite, mono 1-D array"
                )
            waveforms.append(array)
        if not waveforms:
            raise ValueError("The audio batch cannot be empty")
        if max_length is not None and max_length <= 0:
            raise ValueError("max_length must be positive")
        return self.pad(
            BatchFeature({"input_values": waveforms}),
            padding=padding,
            max_length=max_length,
            truncation=truncation,
            pad_to_multiple_of=pad_to_multiple_of,
            return_attention_mask=(
                self.return_attention_mask
                if return_attention_mask is None
                else return_attention_mask
            ),
            return_tensors=return_tensors,
        )


AudioEmbeddingFeatureExtractor.register_for_auto_class("AutoFeatureExtractor")