from __future__ import annotations

import logging
from typing import Any

try:
    from mlx_lm import generate, load

    HAS_MLX = True
except ImportError:
    HAS_MLX = False
    generate = None  # type: ignore
    load = None  # type: ignore

logger = logging.getLogger("dscons.llm.mlx")


class MlxBackendMixin:
    """Local Apple Silicon MLX inference backend mixin."""

    def _ensure_model_loaded(self) -> tuple[Any, Any]:
        """Load and cache the configured MLX model on first use."""
        if not HAS_MLX:
            raise RuntimeError(
                "Thư viện MLX không khả dụng trên môi trường hiện tại (chỉ hỗ trợ macOS Apple Silicon)."
            )
        if self._model is None or self._tokenizer is None:
            self._model, self._tokenizer = load(self._model_path)
        return self._model, self._tokenizer

    def _build_prompt(self, messages: list[dict[str, str]]) -> str:
        """Convert OpenAI-style messages into a prompt string for MLX."""
        _, tokenizer = self._ensure_model_loaded()

        if hasattr(tokenizer, "apply_chat_template"):
            prompt = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
            if isinstance(prompt, str):
                return prompt

        rendered_messages: list[str] = []
        for message in messages:
            role = str(message.get("role", "user")).strip() or "user"
            content = str(message.get("content", ""))
            rendered_messages.append(f"{role.upper()}: {content}")
        rendered_messages.append("ASSISTANT:")
        return "\n\n".join(rendered_messages)

    def _generate_sync_mlx(
        self,
        messages: list[dict[str, str]],
        *,
        model: str | None = None,
        temperature: float = 0.2,
        max_tokens: int | None = None,
        extra_body: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Run synchronous MLX generation and map the result to OpenAI-like output."""
        del extra_body

        if model and model != self._model_path:
            loaded_model, loaded_tokenizer = load(model)
        else:
            loaded_model, loaded_tokenizer = self._ensure_model_loaded()

        prompt = self._build_prompt(messages)
        effective_max_tokens = max_tokens or self._default_max_tokens

        text = generate(
            loaded_model,
            loaded_tokenizer,
            prompt=prompt,
            max_tokens=effective_max_tokens,
            temp=temperature,
            verbose=False,
        )

        if not isinstance(text, str):
            text = str(text)

        return {
            "id": "mlx-chat-completion",
            "object": "chat.completion",
            "model": model or self._model_path,
            "choices": [
                {
                    "index": 0,
                    "finish_reason": "stop",
                    "message": {
                        "role": "assistant",
                        "content": text,
                    },
                }
            ],
        }