"""OpenAI SDK auto-instrumentation.

Patches ``openai.ChatCompletion.create`` (and async variant) to
automatically capture model_call spans with token counts and cost.
"""
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="openai",
    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 OpenAIInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_create: Any = None
        self._original_acreate: Any = None

    def instrument(self) -> None:
        import openai

        completions_cls = openai.resources.chat.completions.Completions
        self._original_create = completions_cls.create
        completions_cls.create = wrap_sync(self._original_create, _CONFIG)

        async_completions_cls = openai.resources.chat.completions.AsyncCompletions
        self._original_acreate = async_completions_cls.create
        async_completions_cls.create = wrap_async(self._original_acreate, _CONFIG)

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

        openai.resources.chat.completions.Completions.create = self._original_create
        openai.resources.chat.completions.AsyncCompletions.create = self._original_acreate
        self._original_create = None
        self._original_acreate = None
