"""Mistral adapter — instruments ``client.chat.complete``.

Recognizes a ``mistralai.Mistral`` client by module. Mistral's response shape
matches OpenAI's chat completions (``choices[].message.content`` and
``usage.total_tokens``), so the response parsing is shared with the OpenAI
adapter; only the recognized module and the wrapped method differ. The
``mistralai`` package is NOT a hard dependency — the adapter duck-types.
"""
from __future__ import annotations

import functools
from typing import Any

from .base import AbomRun, ModelCall, _sha256, register
from .openai import _output_text, _tokens

MISTRAL_ENDPOINT = "api.mistral.ai"


def parse(model: str, response: Any) -> ModelCall:
    """Normalize a Mistral chat response into a :class:`ModelCall`.

    Records the requested model, the total token usage, and a hash of the
    output text (shape is OpenAI-compatible, so the readers are reused).
    """
    text = _output_text(response)
    return ModelCall(
        model=model,
        tokens=_tokens(response),
        output_sha256=_sha256(text) if text else None,
        endpoint=MISTRAL_ENDPOINT,
    )


class MistralAdapter:
    provider = "mistral"

    def matches(self, client: Any) -> bool:
        mod = type(client).__module__ or ""
        name = type(client).__name__
        return mod.split(".")[0] == "mistralai" or name == "Mistral"

    def wrap(self, client: Any, run: AbomRun) -> Any:
        chat = getattr(client, "chat", None)
        if chat is None or not hasattr(chat, "complete"):
            raise TypeError(
                "Mistral adapter expects a client with .chat.complete; "
                f"got {type(client).__name__}"
            )
        original = chat.complete

        @functools.wraps(original)
        def complete(*args: Any, **kwargs: Any) -> Any:
            model = kwargs.get("model") or (args[0] if args else "unknown")
            # gate BEFORE the call leaves the boundary (deny-by-default)
            run.enforce(ModelCall(model=model, endpoint=MISTRAL_ENDPOINT))
            response = original(*args, **kwargs)
            run.record(parse(model, response))
            return response

        chat.complete = complete  # type: ignore[attr-defined]
        return client


register(MistralAdapter())
