"""
MLX adapter — Apple Silicon optimized models.

Supports:
- Loading MLX models from HuggingFace Hub
- Converting HuggingFace models to MLX format
- Wrapping MLX models as PyTorch-compatible for stacking

Note: MLX is only available on Apple Silicon Macs.

Usage::

    from model_osmosis.adapters.mlx import MLXAdapter

    adapter = MLXAdapter()

    # Load an MLX model from Hub
    model = adapter.load("mlx-community/Llama-3-8B-Instruct-4bit")

    # Convert from HuggingFace
    mlx_expert = adapter.from_huggingface(
        "meta-llama/Llama-2-7b-hf",
        quantize=True,
        bits=4,
    )

    # Convert to PyTorch-compatible wrapper
    pytorch_expert = model.to_pytorch()
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np


class MLXExpert:
    """Wrapped MLX model compatible with ModelOsmosis.

    MLX models run on Apple Silicon and can be converted to
    PyTorch wrappers for stacking.

    Args:
        model: The MLX model.
        tokenizer: Associated tokenizer.
        feature_dim: Output feature dimension.
        name: Human-readable name.
        model_path: Path to the model.
    """

    def __init__(
        self,
        model: Any,
        tokenizer: Any,
        feature_dim: int,
        name: str = "mlx_expert",
        model_path: Optional[str] = None,
    ) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.feature_dim = feature_dim
        self.hidden_size = feature_dim
        self.name = name
        self.model_path = model_path

    def __call__(self, x: Any) -> np.ndarray:
        """Forward pass using MLX.

        Args:
            x: Input (can be numpy array, MLX array, or text).

        Returns:
            Output as numpy array.
        """
        try:
            import mlx.core as mx
        except ImportError:
            raise ImportError(
                "MLX is required. Install with: pip install mlx"
            )

        # Convert input if needed
        if isinstance(x, np.ndarray):
            x = mx.array(x)

        # Run model
        outputs = self.model(x)

        # Extract hidden states if available
        if hasattr(outputs, "last_hidden_state"):
            result = outputs.last_hidden_state[:, -1, :]
        elif isinstance(outputs, tuple):
            result = outputs[0]
            if result.ndim == 3:
                result = result[:, -1, :]
        else:
            result = outputs

        # Convert to numpy
        return np.array(result)

    def generate(
        self,
        prompt: str,
        max_tokens: int = 100,
        temperature: float = 0.7,
        **kwargs,
    ) -> str:
        """Generate text from prompt.

        Args:
            prompt: Input prompt.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.

        Returns:
            Generated text.
        """
        try:
            from mlx_lm import generate
        except ImportError:
            raise ImportError(
                "mlx-lm is required for generation. "
                "Install with: pip install mlx-lm"
            )

        return generate(
            self.model,
            self.tokenizer,
            prompt=prompt,
            max_tokens=max_tokens,
            temp=temperature,
            **kwargs,
        )

    def to_pytorch(self) -> "MLXPyTorchWrapper":
        """Convert to PyTorch-compatible wrapper.

        Returns:
            nn.Module wrapper that calls MLX internally.
        """
        import torch
        import torch.nn as nn

        class MLXPyTorchWrapper(nn.Module):
            def __init__(self, mlx_expert: MLXExpert):
                super().__init__()
                self.mlx_expert = mlx_expert
                self.feature_dim = mlx_expert.feature_dim
                self.hidden_size = mlx_expert.hidden_size
                self.name = mlx_expert.name

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                x_np = x.detach().cpu().numpy()
                out_np = self.mlx_expert(x_np)
                return torch.from_numpy(out_np).to(x.device)

        return MLXPyTorchWrapper(self)


class MLXAdapter:
    """Load and convert MLX models.

    Note: Only available on Apple Silicon Macs.

    Args:
        verbose: Print loading progress.
    """

    def __init__(self, verbose: bool = True) -> None:
        self.verbose = verbose

        # Check MLX availability
        try:
            import mlx.core as mx
            self._mlx_available = True
        except ImportError:
            self._mlx_available = False

    def _check_mlx(self) -> None:
        """Check if MLX is available."""
        if not self._mlx_available:
            raise ImportError(
                "MLX is required for Apple Silicon support. "
                "Install with: pip install mlx mlx-lm\n"
                "Note: MLX only works on Apple Silicon Macs."
            )

    def load(
        self,
        model_name: str,
        name: Optional[str] = None,
        adapter_path: Optional[str] = None,
    ) -> MLXExpert:
        """Load an MLX model from HuggingFace Hub.

        Args:
            model_name: HuggingFace model ID (e.g., "mlx-community/Llama-3-8B").
            name: Human-readable name.
            adapter_path: Optional path to LoRA adapter.

        Returns:
            MLXExpert ready for inference.
        """
        self._check_mlx()

        try:
            from mlx_lm import load
        except ImportError:
            raise ImportError(
                "mlx-lm is required. Install with: pip install mlx-lm"
            )

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

        # Load model and tokenizer
        model, tokenizer = load(model_name, adapter_path=adapter_path)

        # Infer feature dim from config
        if hasattr(model, "config"):
            feature_dim = getattr(model.config, "hidden_size", 4096)
        elif hasattr(model, "args"):
            feature_dim = getattr(model.args, "dim", 4096)
        else:
            feature_dim = 4096

        if self.verbose:
            print(f"[MLX] Loaded '{name or model_name.split('/')[-1]}': feature_dim={feature_dim}")

        return MLXExpert(
            model=model,
            tokenizer=tokenizer,
            feature_dim=feature_dim,
            name=name or model_name.split("/")[-1],
            model_path=model_name,
        )

    def from_huggingface(
        self,
        model_name: str,
        output_path: Optional[str] = None,
        quantize: bool = True,
        q_bits: int = 4,
        name: Optional[str] = None,
    ) -> MLXExpert:
        """Convert a HuggingFace model to MLX format.

        Args:
            model_name: HuggingFace model ID.
            output_path: Where to save converted model (temp if None).
            quantize: Whether to quantize the model.
            q_bits: Quantization bits (4 or 8).
            name: Human-readable name.

        Returns:
            MLXExpert with converted model.
        """
        self._check_mlx()

        try:
            from mlx_lm import convert
        except ImportError:
            raise ImportError(
                "mlx-lm is required for conversion. "
                "Install with: pip install mlx-lm"
            )

        import tempfile

        if output_path is None:
            output_path = tempfile.mkdtemp(prefix="mlx_")

        if self.verbose:
            print(f"[MLX] Converting '{model_name}' to MLX format...")
            if quantize:
                print(f"[MLX] Quantizing to {q_bits}-bit...")

        # Convert model
        convert(
            model_name,
            mlx_path=output_path,
            quantize=quantize,
            q_bits=q_bits if quantize else None,
        )

        if self.verbose:
            print(f"[MLX] Saved to: {output_path}")

        # Load the converted model
        return self.load(output_path, name=name)

    def list_mlx_models(self, filter_str: str = "") -> List[str]:
        """List available MLX models on HuggingFace Hub.

        Args:
            filter_str: Optional filter string.

        Returns:
            List of model IDs.
        """
        try:
            from huggingface_hub import list_models
        except ImportError:
            return []

        models = list_models(
            author="mlx-community",
            search=filter_str,
        )

        return [m.id for m in models]
