"""
Audio model adapter — load Whisper, Wav2Vec2, HuBERT, and other audio models.

Supports:
- Whisper (OpenAI) — speech recognition encoder
- Wav2Vec2 (Facebook) — self-supervised audio representations
- HuBERT (Facebook) — hidden unit BERT for audio
- WavLM (Microsoft) — speech pre-training with masked prediction
- Other HuggingFace audio models

Usage::

    from model_osmosis.adapters.audio import AudioAdapter

    adapter = AudioAdapter()

    # Load Whisper encoder
    whisper = adapter.load("openai/whisper-base")

    # Load Wav2Vec2
    wav2vec = adapter.load("facebook/wav2vec2-base-960h")

    # Load HuBERT
    hubert = adapter.load("facebook/hubert-base-ls960")
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional, Union

import torch
import torch.nn as nn
import numpy as np


class AudioModelType(Enum):
    """Types of audio models."""

    WHISPER = "whisper"
    WAV2VEC2 = "wav2vec2"
    HUBERT = "hubert"
    WAVLM = "wavlm"
    GENERIC = "generic"


class AudioExpert(nn.Module):
    """Wrapped audio model as stackable expert.

    Args:
        model: The audio model.
        feature_dim: Output feature dimension.
        name: Human-readable name.
        sample_rate: Expected audio sample rate.
        model_type: Type of audio model.
    """

    def __init__(
        self,
        model: nn.Module,
        feature_dim: int,
        name: str = "audio_expert",
        sample_rate: int = 16000,
        model_type: AudioModelType = AudioModelType.GENERIC,
    ) -> None:
        super().__init__()
        self.model = model
        self.feature_dim = feature_dim
        self.hidden_size = feature_dim  # Alias for compatibility
        self.name = name
        self.sample_rate = sample_rate
        self.model_type = model_type

        # Freeze weights
        for param in self.model.parameters():
            param.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Extract features from audio.

        Args:
            x: Audio tensor:
                - Raw waveform (batch, samples)
                - Mel features (batch, features, time) for Whisper
                - Pre-computed embeddings (batch, seq_len, hidden)

        Returns:
            Feature tensor (batch, feature_dim).
        """
        with torch.no_grad():
            if self.model_type == AudioModelType.WHISPER:
                # Whisper uses encoder for feature extraction
                if hasattr(self.model, "encoder"):
                    outputs = self.model.encoder(x)
                else:
                    outputs = self.model(x)
            elif self.model_type in (AudioModelType.WAV2VEC2, AudioModelType.HUBERT, AudioModelType.WAVLM):
                # These take raw waveform as input_values
                if x.dim() == 2:
                    outputs = self.model(input_values=x)
                else:
                    outputs = self.model(x)
            else:
                # Generic: try common input names
                try:
                    outputs = self.model(input_values=x)
                except TypeError:
                    try:
                        outputs = self.model(input_features=x)
                    except TypeError:
                        outputs = self.model(x)

        # Extract hidden states
        if hasattr(outputs, "last_hidden_state"):
            hidden = outputs.last_hidden_state
        elif isinstance(outputs, tuple):
            hidden = outputs[0]
        else:
            hidden = outputs

        # Pool over time dimension
        if hidden.dim() == 3:
            hidden = hidden.mean(dim=1)

        return hidden

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, feature_dim={self.feature_dim}, "
            f"sample_rate={self.sample_rate}, model_type={self.model_type.value}"
        )


class AudioAdapter:
    """Load audio models from HuggingFace.

    Args:
        device: Device to load models to.
        torch_dtype: Data type for model weights.
        verbose: Print loading progress.
    """

    # Model type detection patterns
    MODEL_PATTERNS = {
        AudioModelType.WHISPER: ["whisper"],
        AudioModelType.WAV2VEC2: ["wav2vec2", "wav2vec"],
        AudioModelType.HUBERT: ["hubert"],
        AudioModelType.WAVLM: ["wavlm"],
    }

    # Feature dimensions for common models
    DIM_LOOKUP = {
        "whisper-tiny": 384,
        "whisper-base": 512,
        "whisper-small": 768,
        "whisper-medium": 1024,
        "whisper-large": 1280,
        "wav2vec2-base": 768,
        "wav2vec2-large": 1024,
        "hubert-base": 768,
        "hubert-large": 1024,
        "hubert-xlarge": 1280,
        "wavlm-base": 768,
        "wavlm-large": 1024,
    }

    def __init__(
        self,
        device: str = "cpu",
        torch_dtype: Optional[torch.dtype] = None,
        verbose: bool = True,
    ) -> None:
        self.device = device
        self.torch_dtype = torch_dtype
        self.verbose = verbose

    def load(
        self,
        model_name: str,
        name: Optional[str] = None,
        **kwargs,
    ) -> AudioExpert:
        """Load an audio model.

        Args:
            model_name: HuggingFace model ID.
            name: Human-readable name.

        Returns:
            AudioExpert ready for stacking.
        """
        try:
            from transformers import AutoConfig
        except ImportError:
            raise ImportError(
                "Audio adapter requires `transformers`. "
                "Install with: pip install transformers"
            )

        if self.verbose:
            print(f"[Audio] Loading '{model_name}'...")

        # Detect model type
        model_type = self._detect_type(model_name)

        # Load config
        config = AutoConfig.from_pretrained(model_name)

        # Get feature dimension
        feature_dim = self._get_feature_dim(config, model_name)

        # Get sample rate
        sample_rate = getattr(config, "sample_rate", 16000)

        # Load model
        model = self._load_model(model_name, model_type, config, kwargs)

        model = model.to(self.device)
        if self.torch_dtype:
            model = model.to(self.torch_dtype)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(
                f"[Audio] Loaded '{name or model_name.split('/')[-1]}': "
                f"{total_params:,} params, feature_dim={feature_dim}, "
                f"sample_rate={sample_rate}"
            )

        return AudioExpert(
            model=model,
            feature_dim=feature_dim,
            name=name or model_name.split("/")[-1],
            sample_rate=sample_rate,
            model_type=model_type,
        )

    def _detect_type(self, model_name: str) -> AudioModelType:
        """Detect audio model type from name."""
        name_lower = model_name.lower()
        for model_type, patterns in self.MODEL_PATTERNS.items():
            for pattern in patterns:
                if pattern in name_lower:
                    return model_type
        return AudioModelType.GENERIC

    def _get_feature_dim(self, config: Any, model_name: str) -> int:
        """Get feature dimension from config or lookup."""
        # Try config attributes
        for attr in ["hidden_size", "d_model", "encoder_hidden_size"]:
            dim = getattr(config, attr, None)
            if dim is not None:
                return dim

        # Try lookup table
        name_lower = model_name.lower()
        for key, dim in self.DIM_LOOKUP.items():
            if key in name_lower:
                return dim

        # Default
        return 768

    def _load_model(
        self,
        model_name: str,
        model_type: AudioModelType,
        config: Any,
        kwargs: Dict,
    ) -> nn.Module:
        """Load the appropriate model class."""
        from transformers import AutoModel

        if model_type == AudioModelType.WHISPER:
            # For Whisper, we want the encoder model
            from transformers import WhisperModel
            model = WhisperModel.from_pretrained(model_name, **kwargs)
        else:
            model = AutoModel.from_pretrained(model_name, **kwargs)

        return model

    def load_multiple(
        self,
        model_names: List[str],
        names: Optional[List[str]] = None,
    ) -> List[AudioExpert]:
        """Load multiple audio models.

        Args:
            model_names: List of HuggingFace model IDs.
            names: Optional list of names.

        Returns:
            List of AudioExpert modules.
        """
        names = names or [None] * len(model_names)
        experts = []
        for model_name, name in zip(model_names, names):
            expert = self.load(model_name, name=name)
            experts.append(expert)
        return experts

    def get_feature_dim(self, model_name: str) -> int:
        """Get feature dimension without loading full model.

        Args:
            model_name: HuggingFace model ID.

        Returns:
            Feature dimension.
        """
        try:
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained(model_name)
            return self._get_feature_dim(config, model_name)
        except Exception:
            # Try lookup
            name_lower = model_name.lower()
            for key, dim in self.DIM_LOOKUP.items():
                if key in name_lower:
                    return dim
            return 768
