"""Mistral AI SDK auto-instrumentation.

Patches ``mistralai.Mistral.chat.complete`` and ``complete_async``
to capture model_call spans.
"""
from __future__ import annotations

from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, wrap_async, wrap_sync

_CONFIG = ProviderConfig(
    name="mistral",
    input_tokens_field="prompt_tokens",
    output_tokens_field="completion_tokens",
    finish_reason_getter=lambda r: getattr(r.choices[0], "finish_reason", None) if getattr(r, "choices", None) else None,
)


class MistralInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_complete: Any = None
        self._original_acomplete: Any = None

    def instrument(self) -> None:
        from mistralai import Mistral

        chat = Mistral.chat
        chat_cls = type(chat) if not isinstance(chat, type) else chat

        # mistralai v1: Mistral().chat.complete / complete_async
        # The chat attribute is a Chat resource with complete/complete_async methods
        try:
            from mistralai.resources import ChatResource
        except ImportError:
            # Fallback: try to find the class from the chat property
            ChatResource = type(Mistral.__new__(Mistral).chat) if hasattr(Mistral, "chat") else None  # type: ignore[assignment]

        if ChatResource is None:
            return

        self._original_complete = ChatResource.complete
        ChatResource.complete = wrap_sync(self._original_complete, _CONFIG)

        self._original_acomplete = ChatResource.complete_async
        ChatResource.complete_async = wrap_async(self._original_acomplete, _CONFIG)

    def uninstrument(self) -> None:
        if self._original_complete is None:
            return

        try:
            from mistralai.resources import ChatResource
        except ImportError:
            return

        ChatResource.complete = self._original_complete
        ChatResource.complete_async = self._original_acomplete
        self._original_complete = None
        self._original_acomplete = None
