"""
HuggingFace adapter — load models from HuggingFace Hub and
wrap them as nn.Module experts for stacking.

Supports multiple modalities:
- Text: LLMs (Llama, Mistral, Qwen, etc.)
- Vision: ViT, ResNet, ConvNeXt, DINOv2, etc.
- Audio: Whisper, Wav2Vec2, HuBERT, etc.
- Multimodal: CLIP, LLaVA, BLIP, etc.

Usage::

    from model_osmosis.adapters.huggingface import HuggingFaceAdapter

    adapter = HuggingFaceAdapter()

    # Load text models
    text_expert = adapter.load("Qwen/Qwen2-0.5B")

    # Load vision models (auto-detected)
    vision_expert = adapter.load("google/vit-base-patch16-224")

    # Load audio models (auto-detected)
    audio_expert = adapter.load("openai/whisper-base")

    # Load multimodal models (auto-detected)
    clip_expert = adapter.load("openai/clip-vit-base-patch32")
"""

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 ModelModality(Enum):
    """Supported model modalities."""

    TEXT = "text"
    VISION = "vision"
    AUDIO = "audio"
    MULTIMODAL = "multimodal"
    UNKNOWN = "unknown"


class HuggingFaceExpert(nn.Module):
    """Wraps a HuggingFace model as an nn.Module expert for stacking.

    This extracts the transformer body (without the LM head) so it can
    be used as an expert in a StackedModel. The input/output are tensors,
    not token IDs.

    Args:
        hf_model: The loaded HuggingFace model.
        hidden_size: The model's hidden dimension.
        name: Human-readable name for this expert.
        modality: The model's modality (text, vision, audio, multimodal).
        use_hidden_states: Whether to use the hidden states output.
    """

    def __init__(
        self,
        hf_model: nn.Module,
        hidden_size: int,
        name: str = "hf_expert",
        modality: ModelModality = ModelModality.TEXT,
        use_hidden_states: bool = True,
    ) -> None:
        super().__init__()
        self.model = hf_model
        self.hidden_size = hidden_size
        self.name = name
        self.modality = modality
        self.use_hidden_states = use_hidden_states

        # Freeze expert weights (they're pretrained, don't train them)
        for param in self.model.parameters():
            param.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Accepts raw tensor input and returns hidden state output.

        Args:
            x: Input tensor of shape (batch, hidden_size).

        Returns:
            Output tensor of shape (batch, hidden_size).
        """
        # If the input is 2D, add a sequence length dimension
        if x.dim() == 2:
            x = x.unsqueeze(1)  # (batch, 1, hidden)

        # Run through the model body
        with torch.no_grad():
            outputs = self.model(inputs_embeds=x)

        # Extract hidden state
        if hasattr(outputs, "last_hidden_state"):
            hidden = outputs.last_hidden_state
        elif hasattr(outputs, "hidden_states") and outputs.hidden_states:
            hidden = outputs.hidden_states[-1]
        else:
            # Fallback: treat output as tensor
            hidden = outputs[0] if isinstance(outputs, tuple) else outputs

        # Pool: take the last token's hidden state
        if hidden.dim() == 3:
            hidden = hidden[:, -1, :]  # (batch, hidden_size)

        return hidden

    def extra_repr(self) -> str:
        return f"name={self.name}, hidden_size={self.hidden_size}, modality={self.modality.value}"


class VisionExpert(HuggingFaceExpert):
    """Specialized expert wrapper for vision models (ViT, ResNet, etc.).

    Handles:
    - Pixel values input (batch, channels, height, width)
    - Pre-computed embeddings (batch, seq_len, hidden)

    Args:
        hf_model: The loaded HuggingFace vision model.
        hidden_size: The model's hidden dimension.
        name: Human-readable name for this expert.
        image_size: Expected input image size.
        patch_size: Patch size for ViT-style models.
    """

    def __init__(
        self,
        hf_model: nn.Module,
        hidden_size: int,
        name: str = "vision_expert",
        image_size: int = 224,
        patch_size: int = 16,
    ) -> None:
        super().__init__(hf_model, hidden_size, name, ModelModality.VISION)
        self.image_size = image_size
        self.patch_size = patch_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass for vision models.

        Args:
            x: Either:
                - Pixel values (batch, channels, height, width)
                - Pre-computed embeddings (batch, seq_len, hidden) or (batch, hidden)

        Returns:
            Feature tensor of shape (batch, hidden_size).
        """
        with torch.no_grad():
            if x.dim() == 4:
                # Image input: (batch, channels, height, width)
                outputs = self.model(pixel_values=x)
            elif x.dim() == 3:
                # Embedding input: (batch, seq_len, hidden)
                outputs = self.model(inputs_embeds=x)
            else:
                # 2D input: (batch, hidden) -> add seq dim
                x = x.unsqueeze(1)
                outputs = self.model(inputs_embeds=x)

        # Extract pooled or CLS token representation
        if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
            return outputs.pooler_output
        elif hasattr(outputs, "last_hidden_state"):
            return outputs.last_hidden_state[:, 0, :]  # CLS token
        elif isinstance(outputs, tuple):
            hidden = outputs[0]
            if hidden.dim() == 3:
                return hidden[:, 0, :]
            return hidden
        return outputs

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, hidden_size={self.hidden_size}, "
            f"image_size={self.image_size}, patch_size={self.patch_size}"
        )


class AudioExpert(HuggingFaceExpert):
    """Specialized expert wrapper for audio models (Whisper, Wav2Vec2, etc.).

    Handles:
    - Raw waveform input (batch, samples)
    - Mel spectrogram features (batch, features, time)

    Args:
        hf_model: The loaded HuggingFace audio model.
        hidden_size: The model's hidden dimension.
        name: Human-readable name for this expert.
        sample_rate: Expected audio sample rate.
        model_type: Type of audio model ('whisper', 'wav2vec2', 'hubert').
    """

    def __init__(
        self,
        hf_model: nn.Module,
        hidden_size: int,
        name: str = "audio_expert",
        sample_rate: int = 16000,
        model_type: str = "generic",
    ) -> None:
        super().__init__(hf_model, hidden_size, name, ModelModality.AUDIO)
        self.sample_rate = sample_rate
        self.model_type = model_type

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass for audio models.

        Args:
            x: Either:
                - Raw waveform (batch, samples)
                - Mel features (batch, features, time)
                - Pre-computed embeddings (batch, seq_len, hidden)

        Returns:
            Feature tensor of shape (batch, hidden_size).
        """
        with torch.no_grad():
            if self.model_type == "whisper":
                # Whisper uses encoder only for feature extraction
                if hasattr(self.model, "encoder"):
                    outputs = self.model.encoder(x)
                else:
                    outputs = self.model(x)
            elif self.model_type in ("wav2vec2", "hubert"):
                # Wav2Vec2/HuBERT take raw waveform
                outputs = self.model(input_values=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 and pool hidden states
        if hasattr(outputs, "last_hidden_state"):
            hidden = outputs.last_hidden_state
            return hidden.mean(dim=1)  # Pool over time
        elif isinstance(outputs, tuple):
            hidden = outputs[0]
            if hidden.dim() == 3:
                return hidden.mean(dim=1)
            return hidden
        return outputs

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


class MultimodalExpert(HuggingFaceExpert):
    """Specialized expert wrapper for multimodal models (CLIP, LLaVA, etc.).

    Supports encoding from multiple modalities into a shared space.

    Args:
        hf_model: The loaded HuggingFace multimodal model.
        hidden_size: The model's hidden dimension.
        name: Human-readable name for this expert.
        text_dim: Hidden dimension for text features.
        vision_dim: Hidden dimension for vision features.
        model_type: Type of multimodal model ('clip', 'llava', 'blip').
    """

    def __init__(
        self,
        hf_model: nn.Module,
        hidden_size: int,
        name: str = "multimodal_expert",
        text_dim: Optional[int] = None,
        vision_dim: Optional[int] = None,
        model_type: str = "clip",
    ) -> None:
        super().__init__(hf_model, hidden_size, name, ModelModality.MULTIMODAL)
        self.text_dim = text_dim or hidden_size
        self.vision_dim = vision_dim or hidden_size
        self.model_type = model_type

    def encode_text(
        self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Encode text to feature space.

        Args:
            input_ids: Token IDs (batch, seq_len).
            attention_mask: Optional attention mask.

        Returns:
            Text features (batch, text_dim).
        """
        with torch.no_grad():
            if self.model_type == "clip":
                kwargs = {"input_ids": input_ids}
                if attention_mask is not None:
                    kwargs["attention_mask"] = attention_mask
                return self.model.get_text_features(**kwargs)
            else:
                outputs = self.model.text_model(
                    input_ids=input_ids, attention_mask=attention_mask
                )
                if hasattr(outputs, "pooler_output"):
                    return outputs.pooler_output
                return outputs.last_hidden_state[:, 0, :]

    def encode_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Encode images to feature space.

        Args:
            pixel_values: Image tensor (batch, channels, height, width).

        Returns:
            Image features (batch, vision_dim).
        """
        with torch.no_grad():
            if self.model_type == "clip":
                return self.model.get_image_features(pixel_values=pixel_values)
            else:
                outputs = self.model.vision_model(pixel_values=pixel_values)
                if hasattr(outputs, "pooler_output"):
                    return outputs.pooler_output
                return outputs.last_hidden_state[:, 0, :]

    def forward(
        self,
        x: torch.Tensor,
        modality: str = "vision",
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass for specified modality.

        Args:
            x: Input tensor (format depends on modality).
            modality: Which modality to use ('text' or 'vision').
            attention_mask: Optional attention mask for text.

        Returns:
            Feature tensor of shape (batch, hidden_size).
        """
        if modality == "text":
            return self.encode_text(x, attention_mask)
        else:
            return self.encode_image(x)

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, hidden_size={self.hidden_size}, "
            f"text_dim={self.text_dim}, vision_dim={self.vision_dim}, "
            f"model_type={self.model_type}"
        )


class HuggingFaceAdapter:
    """Load HuggingFace models and wrap them as stackable experts.

    Automatically detects model modality (text, vision, audio, multimodal)
    and wraps with the appropriate expert class.

    Args:
        device: Device to load models to ('cpu', 'cuda', 'auto').
        torch_dtype: Data type for model weights (torch.float16, etc.).
        trust_remote_code: Allow running model-specific code from HF Hub.
        verbose: Print loading progress.
    """

    # Architecture patterns for modality detection
    VISION_ARCHITECTURES = [
        "ViTModel", "ViTForImageClassification",
        "ResNetModel", "ResNetForImageClassification",
        "ConvNextModel", "ConvNextForImageClassification",
        "Dinov2Model", "Dinov2ForImageClassification",
        "SwinModel", "SwinForImageClassification",
        "BeitModel", "BeitForImageClassification",
        "DeiTModel", "DeiTForImageClassification",
        "CLIPVisionModel",
    ]

    AUDIO_ARCHITECTURES = [
        "WhisperModel", "WhisperForConditionalGeneration", "WhisperEncoder",
        "Wav2Vec2Model", "Wav2Vec2ForCTC", "Wav2Vec2ForSequenceClassification",
        "HubertModel", "HubertForCTC", "HubertForSequenceClassification",
        "WavLMModel", "WavLMForCTC",
        "SEWModel", "SEWForCTC",
        "UniSpeechModel", "UniSpeechForCTC",
    ]

    MULTIMODAL_ARCHITECTURES = [
        "CLIPModel", "CLIPTextModel",
        "LlavaForConditionalGeneration", "LlavaNextForConditionalGeneration",
        "BlipModel", "BlipForConditionalGeneration",
        "Blip2Model", "Blip2ForConditionalGeneration",
        "GitModel", "GitForCausalLM",
        "Pix2StructModel", "Pix2StructForConditionalGeneration",
        "InstructBlipModel",
    ]

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

        # Check that transformers is available
        try:
            import transformers
            self._transformers = transformers
        except ImportError:
            raise ImportError(
                "HuggingFace adapter requires the `transformers` library. "
                "Install it with: pip install transformers"
            )

    def detect_modality(self, config: Any) -> ModelModality:
        """Detect model modality from config.

        Args:
            config: HuggingFace model config.

        Returns:
            Detected ModelModality.
        """
        architectures = getattr(config, "architectures", []) or []
        arch_str = " ".join(architectures)

        # Check multimodal first (most specific)
        for pattern in self.MULTIMODAL_ARCHITECTURES:
            if pattern in arch_str:
                return ModelModality.MULTIMODAL

        # Check audio
        for pattern in self.AUDIO_ARCHITECTURES:
            if pattern in arch_str:
                return ModelModality.AUDIO

        # Check vision
        for pattern in self.VISION_ARCHITECTURES:
            if pattern in arch_str:
                return ModelModality.VISION

        # Check config attributes for hints
        if hasattr(config, "image_size") or hasattr(config, "patch_size"):
            return ModelModality.VISION
        if hasattr(config, "num_mel_bins") or hasattr(config, "feature_size"):
            return ModelModality.AUDIO

        # Default to text
        return ModelModality.TEXT

    def _get_hidden_size(self, config: Any) -> int:
        """Extract hidden size from config."""
        for attr in ["hidden_size", "d_model", "n_embd", "projection_dim"]:
            val = getattr(config, attr, None)
            if val is not None:
                return val

        # Check nested configs for multimodal
        if hasattr(config, "text_config"):
            text_hidden = getattr(config.text_config, "hidden_size", None)
            if text_hidden:
                return text_hidden
        if hasattr(config, "vision_config"):
            vision_hidden = getattr(config.vision_config, "hidden_size", None)
            if vision_hidden:
                return vision_hidden

        raise ValueError(f"Could not determine hidden_size from config")

    def load(
        self,
        model_name: str,
        name: Optional[str] = None,
        revision: Optional[str] = None,
        modality: Optional[ModelModality] = None,
    ) -> HuggingFaceExpert:
        """Load a HuggingFace model and wrap it as an expert.

        Automatically detects modality unless explicitly specified.

        Args:
            model_name: HuggingFace model ID (e.g., 'google/vit-base-patch16-224').
            name: Optional human-readable name. Defaults to model_name.
            revision: Specific model revision/branch.
            modality: Force a specific modality (auto-detected if None).

        Returns:
            An appropriate Expert subclass (nn.Module) ready for stacking.
        """
        from transformers import AutoModel, AutoConfig

        name = name or model_name.split("/")[-1]

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

        # Load config first to get architecture info
        config = AutoConfig.from_pretrained(
            model_name,
            trust_remote_code=self.trust_remote_code,
            revision=revision,
        )

        # Detect modality
        detected_modality = modality or self.detect_modality(config)

        if self.verbose:
            print(f"[HFAdapter] Detected modality: {detected_modality.value}")

        # Get hidden size
        try:
            hidden_size = self._get_hidden_size(config)
        except ValueError:
            raise ValueError(
                f"Could not determine hidden_size for '{model_name}'. "
                f"Config keys: {list(config.__dict__.keys())}"
            )

        # Load the model based on modality
        load_kwargs = {
            "trust_remote_code": self.trust_remote_code,
        }
        if self.torch_dtype:
            load_kwargs["torch_dtype"] = self.torch_dtype
        if revision:
            load_kwargs["revision"] = revision

        # Load appropriate model class
        if detected_modality == ModelModality.MULTIMODAL:
            model = AutoModel.from_pretrained(model_name, **load_kwargs)
            expert_class = MultimodalExpert
            extra_kwargs = self._get_multimodal_kwargs(config)
        elif detected_modality == ModelModality.VISION:
            model = AutoModel.from_pretrained(model_name, **load_kwargs)
            expert_class = VisionExpert
            extra_kwargs = self._get_vision_kwargs(config)
        elif detected_modality == ModelModality.AUDIO:
            model = self._load_audio_model(model_name, config, load_kwargs)
            expert_class = AudioExpert
            extra_kwargs = self._get_audio_kwargs(config, model_name)
        else:
            model = AutoModel.from_pretrained(model_name, **load_kwargs)
            expert_class = HuggingFaceExpert
            extra_kwargs = {"modality": ModelModality.TEXT}

        model = model.to(self.device)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(
                f"[HFAdapter] Loaded '{name}': {total_params:,} params, "
                f"hidden_size={hidden_size}, modality={detected_modality.value}"
            )

        expert = expert_class(
            hf_model=model,
            hidden_size=hidden_size,
            name=name,
            **extra_kwargs,
        )

        return expert

    def _load_audio_model(
        self, model_name: str, config: Any, load_kwargs: Dict
    ) -> nn.Module:
        """Load an audio model with appropriate class."""
        from transformers import AutoModel

        # For Whisper, we want just the encoder
        if "whisper" in model_name.lower():
            from transformers import WhisperModel
            model = WhisperModel.from_pretrained(model_name, **load_kwargs)
        else:
            model = AutoModel.from_pretrained(model_name, **load_kwargs)

        return model

    def _get_vision_kwargs(self, config: Any) -> Dict[str, Any]:
        """Extract vision-specific kwargs from config."""
        return {
            "image_size": getattr(config, "image_size", 224),
            "patch_size": getattr(config, "patch_size", 16),
        }

    def _get_audio_kwargs(self, config: Any, model_name: str) -> Dict[str, Any]:
        """Extract audio-specific kwargs from config."""
        # Detect audio model type
        model_type = "generic"
        name_lower = model_name.lower()
        if "whisper" in name_lower:
            model_type = "whisper"
        elif "wav2vec" in name_lower:
            model_type = "wav2vec2"
        elif "hubert" in name_lower:
            model_type = "hubert"

        return {
            "sample_rate": getattr(config, "sample_rate", 16000),
            "model_type": model_type,
        }

    def _get_multimodal_kwargs(self, config: Any) -> Dict[str, Any]:
        """Extract multimodal-specific kwargs from config."""
        text_dim = None
        vision_dim = None

        if hasattr(config, "text_config"):
            text_dim = getattr(config.text_config, "hidden_size", None)
        if hasattr(config, "vision_config"):
            vision_dim = getattr(config.vision_config, "hidden_size", None)

        # Detect model type
        model_type = "generic"
        architectures = getattr(config, "architectures", []) or []
        arch_str = " ".join(architectures).lower()
        if "clip" in arch_str:
            model_type = "clip"
        elif "llava" in arch_str:
            model_type = "llava"
        elif "blip" in arch_str:
            model_type = "blip"

        return {
            "text_dim": text_dim,
            "vision_dim": vision_dim,
            "model_type": model_type,
        }

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

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

        Returns:
            List of Expert 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_hidden_size(self, model_name: str) -> int:
        """Get the hidden size of a HuggingFace model without loading it.

        Args:
            model_name: HuggingFace model ID.

        Returns:
            The hidden dimension size.
        """
        from transformers import AutoConfig

        config = AutoConfig.from_pretrained(
            model_name,
            trust_remote_code=self.trust_remote_code,
        )
        return self._get_hidden_size(config)

    def get_modality(self, model_name: str) -> ModelModality:
        """Get the modality of a HuggingFace model without loading it.

        Args:
            model_name: HuggingFace model ID.

        Returns:
            The detected ModelModality.
        """
        from transformers import AutoConfig

        config = AutoConfig.from_pretrained(
            model_name,
            trust_remote_code=self.trust_remote_code,
        )
        return self.detect_modality(config)
