"""
MLX model adapter — Apple Silicon only.

Uses mlx-lm for inference instead of transformers/PyTorch.
Benefits on M-series Macs:
  - Runs on the Neural Engine / GPU via Metal — no CPU saturation
  - 4-bit quantised models fit in unified memory (8 GB usable for 7B)
  - ~3-10x faster than transformers on CPU

Usage:
    from sichgate.models.mlx_adapter import MLXModelAdapter, is_apple_silicon
    if is_apple_silicon():
        adapter = MLXModelAdapter("Qwen/Qwen2.5-1.5B-Instruct")
        response = adapter.query("Hello, who are you?")
"""

from __future__ import annotations

import platform
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ── mlx-community quantised mirrors ──────────────────────────────────────────
# Maps canonical HuggingFace model IDs to their 4-bit mlx-community equivalents.
# mlx-community distributes pre-quantised .safetensors in MLX format so no
# local quantisation step is required.

MLX_REGISTRY: dict[str, str] = {
    # Qwen2 / Qwen2.5
    "Qwen/Qwen2-1.5B-Instruct":           "mlx-community/Qwen2-1.5B-Instruct-4bit",
    "Qwen/Qwen2.5-1.5B-Instruct":         "mlx-community/Qwen2.5-1.5B-Instruct-4bit",
    "Qwen/Qwen2.5-3B-Instruct":           "mlx-community/Qwen2.5-3B-Instruct-4bit",
    "Qwen/Qwen2.5-7B-Instruct":           "mlx-community/Qwen2.5-7B-Instruct-4bit",
    # TinyLlama
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0": "mlx-community/TinyLlama-1.1B-Chat-v1.0-4bit",
    # Medical / biomedical
    "epfl-llm/meditron-7b":               "mlx-community/meditron-7b-4bit",
    "BioMistral/BioMistral-7B":           "mlx-community/BioMistral-7B-4bit",
    # Mistral
    "mistralai/Mistral-7B-Instruct-v0.2": "mlx-community/Mistral-7B-Instruct-v0.2-4bit",
    "mistralai/Mistral-7B-Instruct-v0.3": "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
    # Microsoft Phi
    "microsoft/phi-2":                     "mlx-community/phi-2-hf-4bit-mlx",
    "microsoft/Phi-3-mini-4k-instruct":   "mlx-community/Phi-3-mini-4k-instruct-4bit",
    "microsoft/Phi-3.5-mini-instruct":    "mlx-community/Phi-3.5-mini-instruct-4bit",
    # Llama 3
    "meta-llama/Llama-3.2-1B-Instruct":   "mlx-community/Llama-3.2-1B-Instruct-4bit",
    "meta-llama/Llama-3.2-3B-Instruct":   "mlx-community/Llama-3.2-3B-Instruct-4bit",
}


def is_apple_silicon() -> bool:
    """Return True when running on an Apple Silicon (arm64) Mac."""
    return platform.system() == "Darwin" and platform.machine() == "arm64"


def resolve_mlx_model(model_id: str) -> str:
    """
    Map a HuggingFace model ID to its mlx-community 4-bit mirror.
    Falls back to the original ID if no mapping exists — mlx-lm can load
    standard HF repos directly (it will quantise on first load).
    """
    resolved = MLX_REGISTRY.get(model_id, model_id)
    if resolved != model_id:
        logger.info(f"MLX: {model_id!r} → {resolved!r} (4-bit quantised)")
    return resolved


def mlx_available() -> bool:
    """Return True if mlx-lm is installed."""
    try:
        import mlx_lm  # noqa: F401
        return True
    except ImportError:
        return False


class MLXModelAdapter:
    """
    Wraps mlx_lm for use as a SichGate model backend.

    Lazy-loads on first query so import time stays low.
    Handles both single-turn and multi-turn (conversation_history) queries.

    Args:
        model_id: HuggingFace model ID or local path.  Auto-resolved to the
                  mlx-community 4-bit mirror if a mapping exists.
        max_tokens: Maximum new tokens per response.
    """

    pipeline_type = "text-generation"

    def __init__(self, model_id: str, max_tokens: int = 256):
        self.model_id = model_id
        self.mlx_model_id = resolve_mlx_model(model_id)
        self.max_tokens = max_tokens
        self._model = None
        self._tokenizer = None

    # ── Lazy load ─────────────────────────────────────────────────────────────

    def _load(self) -> None:
        if self._model is not None:
            return
        try:
            from mlx_lm import load
        except ImportError:
            raise RuntimeError(
                "mlx-lm is not installed.  Install it with:\n"
                "  pip install mlx-lm\n"
                "mlx-lm requires Apple Silicon (M1/M2/M3/M4)."
            )
        logger.info(f"MLX: loading {self.mlx_model_id!r} ...")
        self._model, self._tokenizer = load(self.mlx_model_id)
        logger.info(f"MLX: model ready")

    # ── Prompt formatting ─────────────────────────────────────────────────────

    def _format_prompt(
        self,
        prompt: str,
        conversation_history: Optional[list[dict]] = None,
    ) -> str:
        """
        Apply the model's chat template when history is present, otherwise
        return the raw prompt.  Falls back to simple concatenation if the
        tokenizer does not expose apply_chat_template.
        """
        self._load()

        if not conversation_history:
            # Try wrapping in a single-turn chat template (improves instruction following)
            try:
                messages = [{"role": "user", "content": prompt}]
                return self._tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            except Exception:
                return prompt

        # Multi-turn: build full history + new user turn
        messages = conversation_history + [{"role": "user", "content": prompt}]
        try:
            return self._tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        except Exception:
            # Fallback: simple turn concatenation
            parts = []
            for turn in messages:
                role = turn["role"].capitalize()
                parts.append(f"{role}: {turn['content']}")
            return "\n".join(parts) + "\nAssistant:"

    # ── Public interface ──────────────────────────────────────────────────────

    def query(
        self,
        prompt: str,
        conversation_history: Optional[list[dict]] = None,
    ) -> str:
        """Run inference and return the model's response as a plain string."""
        self._load()
        try:
            from mlx_lm import generate
        except ImportError:
            raise RuntimeError("mlx-lm is not installed.  pip install mlx-lm")

        formatted = self._format_prompt(prompt, conversation_history)
        response: str = generate(
            self._model,
            self._tokenizer,
            prompt=formatted,
            max_tokens=self.max_tokens,
            verbose=False,
        )
        return response.strip() or "[No output]"

    def is_reachable(self) -> tuple[bool, str]:
        try:
            self._load()
            return True, "ok"
        except Exception as e:
            return False, str(e)
