"""Anthropic SDK auto-instrumentation.

Patches ``anthropic.Anthropic().messages.create`` 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="anthropic",
    input_tokens_field="input_tokens",
    output_tokens_field="output_tokens",
    finish_reason_getter=lambda r: getattr(r, "stop_reason", None),
)


class AnthropicInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_create: Any = None
        self._original_acreate: Any = None

    def instrument(self) -> None:
        import anthropic

        messages_cls = anthropic.resources.messages.Messages
        self._original_create = messages_cls.create
        messages_cls.create = wrap_sync(self._original_create, _CONFIG)

        async_messages_cls = anthropic.resources.messages.AsyncMessages
        self._original_acreate = async_messages_cls.create
        async_messages_cls.create = wrap_async(self._original_acreate, _CONFIG)

    def uninstrument(self) -> None:
        if self._original_create is None:
            return
        import anthropic

        anthropic.resources.messages.Messages.create = self._original_create
        anthropic.resources.messages.AsyncMessages.create = self._original_acreate
        self._original_create = None
        self._original_acreate = None
