from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

import torch
from tqdm.auto import tqdm
from transformers import AutoProcessor, EncodecModel

from mteb._create_dataloaders import AudioCollator
from mteb._requires_package import requires_audio_dependencies
from mteb.models import ModelMeta
from mteb.models.abs_encoder import AbsEncoder

if TYPE_CHECKING:
    from torch.utils.data import DataLoader

    from mteb import TaskMetadata
    from mteb.types import Array, BatchedInput, PromptType
    from mteb.types._encoder_io import AudioInput


logger = logging.getLogger(__name__)


class EncodecWrapper(AbsEncoder):
    def __init__(
        self,
        model_name: str,
        revision: str,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        max_audio_length_seconds: float = 30.0,
        **kwargs: Any,
    ):
        requires_audio_dependencies()
        self.model_name = model_name
        self.device = device
        self.max_audio_length_seconds = max_audio_length_seconds

        self.model = EncodecModel.from_pretrained(model_name, revision=revision).to(
            device
        )
        self.model.eval()

        self.processor = AutoProcessor.from_pretrained(model_name)
        self.sampling_rate = self.processor.sampling_rate  # 24000 Hz typically

    def get_audio_embeddings(
        self,
        inputs: DataLoader[AudioInput],
        show_progress_bar: bool = True,
        **kwargs: Any,
    ) -> Array:
        max_samples = int(self.max_audio_length_seconds * self.sampling_rate)
        inputs.collate_fn = AudioCollator(self.sampling_rate, max_samples)

        all_embeddings = []

        for batch in tqdm(
            inputs,
            disable=not show_progress_bar,
        ):
            audio_array = [audio["array"] for audio in batch["audio"]]
            audio_arrays = []
            for array in audio_array:
                # Ensure minimum length for encoder (Encodec needs ~320 samples per frame)
                # Use 1 second minimum to be safe
                min_samples = self.sampling_rate
                if array.shape[-1] < min_samples:
                    padding = torch.zeros(min_samples - array.shape[-1])
                    array = torch.cat([array, padding])

                audio_arrays.append(array.numpy())

            with torch.no_grad():
                # Use processor for batch padding (truncation/min-length done manually above)
                processed = self.processor(
                    raw_audio=audio_arrays,
                    sampling_rate=self.sampling_rate,
                    padding=True,
                    return_tensors="pt",
                )
                input_values = processed["input_values"].to(self.device)

                # Add channel dimension if needed (B, T) -> (B, 1, T)
                if input_values.dim() == 2:
                    input_values = input_values.unsqueeze(1)

                # Get the latent representations directly from the encoder
                latent = self.model.encoder(input_values)

                # Validate latent has time frames
                if latent.shape[2] == 0:
                    raise ValueError(
                        f"Encodec encoder produced 0 time frames. "
                        f"Input shape: {input_values.shape}, latent shape: {latent.shape}"
                    )

                # Apply mean pooling over the time dimension to get fixed-size embeddings
                embeddings = torch.mean(latent, dim=2)  # Average over time dimension

                # Normalize embeddings
                embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)

                all_embeddings.append(embeddings.cpu().detach())

        return torch.cat(all_embeddings, dim=0).numpy()

    def encode(
        self,
        inputs: DataLoader[BatchedInput],
        *,
        task_metadata: TaskMetadata,
        hf_split: str,
        hf_subset: str,
        prompt_type: PromptType | None = None,
        **kwargs: Any,
    ) -> Array:
        if "audio" not in inputs.dataset.features:
            raise ValueError("EncodecWrapper only supports audio inputs.")
        return self.get_audio_embeddings(inputs, **kwargs)


encodec_24khz = ModelMeta(
    loader=EncodecWrapper,
    name="facebook/encodec_24khz",
    languages=["eng-Latn"],
    open_weights=True,
    revision="c1dbe2ae3f1de713481a3b3e7c47f357092ee040",
    release_date="2022-10-25",
    max_tokens=None,
    n_parameters=23_273_218,
    n_embedding_parameters=0,
    memory_usage_mb=88,
    embed_dim=128,
    license="cc-by-nc-4.0",
    reference="https://huggingface.co/facebook/encodec_24khz",
    similarity_fn_name="cosine",
    framework=["PyTorch"],
    use_instructions=False,
    public_training_code="https://github.com/facebookresearch/encodec",
    public_training_data=None,
    training_datasets=None,  # ["AudioSet", "VCTK", "DNS-Challenge"],
    modalities=["audio"],
    citation="""
@misc{défossez2022highfidelityneuralaudio,
      title={High Fidelity Neural Audio Compression},
      author={Alexandre Défossez and Jade Copet and Gabriel Synnaeve and Yossi Adi},
      year={2022},
      eprint={2210.13438},
      archivePrefix={arXiv},
      primaryClass={eess.AS},
      url={https://arxiv.org/abs/2210.13438},
}""",
)
