"""Tests for LiteLLM auto-instrumentation."""
from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from fibinaci._testing.mock_client import MockFibinaci
from fibinaci.instrumentation._wrapper import ProviderConfig, _extract_span


@dataclass
class FakeUsage:
    prompt_tokens: int = 90
    completion_tokens: int = 35
    total_tokens: int = 125


@dataclass
class FakeChoice:
    finish_reason: str = "stop"
    message: Any = None


@dataclass
class FakeResponse:
    model: str = "gpt-4o"
    usage: FakeUsage = None
    choices: list = None

    def __post_init__(self):
        if self.usage is None:
            self.usage = FakeUsage()
        if self.choices is None:
            self.choices = [FakeChoice()]


class TestLiteLLMInstrumentation:
    @pytest.fixture
    def setup_litellm(self):
        mod = MagicMock()
        original_completion = MagicMock(return_value=FakeResponse())
        original_acompletion = MagicMock(return_value=FakeResponse())
        mod.completion = original_completion
        mod.acompletion = original_acompletion
        return mod, original_completion, original_acompletion

    def test_instrument_patches_completion(self, setup_litellm, mock_fibinaci: MockFibinaci):
        mod, original_completion, _ = setup_litellm
        with patch.dict("sys.modules", {"litellm": mod}):
            from fibinaci.instrumentation.litellm import LiteLLMInstrumentor
            inst = LiteLLMInstrumentor()
            inst.instrument()
            assert mod.completion is not original_completion
            inst.uninstrument()

    def test_span_captured_on_success(self, mock_fibinaci: MockFibinaci):
        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,
        )
        response = FakeResponse()
        started = datetime(2025, 3, 29, 12, 0, 0, tzinfo=timezone.utc)
        ended = datetime(2025, 3, 29, 12, 0, 2, tzinfo=timezone.utc)

        span_dict = _extract_span(config, {"model": "gpt-4o"}, response, started, ended)

        assert span_dict["span_kind"] == "model_call"
        assert span_dict["provider"] == "litellm"
        assert span_dict["tokens_input"] == 90
        assert span_dict["tokens_output"] == 35
        assert span_dict["status"] == "ok"
        assert span_dict["duration_ms"] == 2000.0
