"""LiteLLM auto-instrumentation.

Patches ``litellm.completion`` and ``litellm.acompletion`` to capture
model_call spans with token counts and cost.
"""
from __future__ import annotations

import functools
import logging
from datetime import datetime, timezone
from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, _extract_span, _make_error_span

logger = logging.getLogger("fibinaci.instrumentation")

_CONFIG = ProviderConfig(
    name="litellm",
    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,
)


def _wrap_completion(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(*args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_extract_span(_CONFIG, kwargs, response, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture litellm span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, kwargs, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture litellm error span", exc_info=True)
            raise

    return wrapper


def _wrap_acompletion(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(*args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_extract_span(_CONFIG, kwargs, response, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async litellm span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                get_buffer().push(_make_error_span(_CONFIG, kwargs, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async litellm error span", exc_info=True)
            raise

    return wrapper


class LiteLLMInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._original_completion: Any = None
        self._original_acompletion: Any = None

    def instrument(self) -> None:
        import litellm

        self._original_completion = litellm.completion
        litellm.completion = _wrap_completion(self._original_completion)

        self._original_acompletion = litellm.acompletion
        litellm.acompletion = _wrap_acompletion(self._original_acompletion)

    def uninstrument(self) -> None:
        if self._original_completion is None:
            return
        import litellm

        litellm.completion = self._original_completion
        litellm.acompletion = self._original_acompletion
        self._original_completion = None
        self._original_acompletion = None
