"""Cohere SDK auto-instrumentation.

Patches ``cohere.Client.chat`` and ``cohere.AsyncClient.chat`` as well as
``generate`` / ``agenerate`` 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="cohere",
    input_tokens_field="input_tokens",
    output_tokens_field="output_tokens",
    finish_reason_getter=lambda r: getattr(r, "finish_reason", None),
)

class CohereInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._originals: dict[str, Any] = {}

    def instrument(self) -> None:
        import cohere

        # Chat
        client_cls = cohere.Client
        if hasattr(client_cls, "chat"):
            self._originals["chat"] = client_cls.chat
            client_cls.chat = wrap_sync(client_cls.chat, _CONFIG)

        # Generate
        if hasattr(client_cls, "generate"):
            self._originals["generate"] = client_cls.generate
            client_cls.generate = wrap_sync(client_cls.generate, _CONFIG)

        # Async variants
        async_cls = cohere.AsyncClient
        if hasattr(async_cls, "chat"):
            self._originals["achat"] = async_cls.chat
            async_cls.chat = wrap_async(async_cls.chat, _CONFIG)

        if hasattr(async_cls, "generate"):
            self._originals["agenerate"] = async_cls.generate
            async_cls.generate = wrap_async(async_cls.generate, _CONFIG)

    def uninstrument(self) -> None:
        if not self._originals:
            return
        import cohere

        client_cls = cohere.Client
        async_cls = cohere.AsyncClient

        if "chat" in self._originals:
            client_cls.chat = self._originals["chat"]
        if "generate" in self._originals:
            client_cls.generate = self._originals["generate"]
        if "achat" in self._originals:
            async_cls.chat = self._originals["achat"]
        if "agenerate" in self._originals:
            async_cls.generate = self._originals["agenerate"]
        self._originals.clear()
