"""Tests for Groq 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 = 110
    completion_tokens: int = 55
    total_tokens: int = 165


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


@dataclass
class FakeResponse:
    model: str = "llama-3.1-70b-versatile"
    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 TestGroqInstrumentation:
    @pytest.fixture
    def setup_groq(self):
        mod = MagicMock()
        completions_cls = MagicMock()
        original_create = MagicMock(return_value=FakeResponse())
        completions_cls.create = original_create
        mod.resources.chat.completions.Completions = completions_cls
        mod.resources.chat.completions.AsyncCompletions = MagicMock()
        return mod, completions_cls, original_create

    def test_instrument_patches_create(self, setup_groq, mock_fibinaci: MockFibinaci):
        mod, completions_cls, original_create = setup_groq
        with patch.dict("sys.modules", {
            "groq": mod,
            "groq.resources": mod.resources,
            "groq.resources.chat": mod.resources.chat,
            "groq.resources.chat.completions": mod.resources.chat.completions,
        }):
            from fibinaci.instrumentation.groq import GroqInstrumentor
            inst = GroqInstrumentor()
            inst.instrument()
            assert completions_cls.create is not original_create
            inst.uninstrument()

    def test_span_captured_on_success(self, mock_fibinaci: MockFibinaci):
        config = ProviderConfig(
            name="groq",
            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, 0, 500000, tzinfo=timezone.utc)

        span_dict = _extract_span(config, {"model": "llama-3.1-70b-versatile"}, response, started, ended)

        assert span_dict["span_kind"] == "model_call"
        assert span_dict["provider"] == "groq"
        assert span_dict["tokens_input"] == 110
        assert span_dict["tokens_output"] == 55
        assert span_dict["status"] == "ok"
        assert span_dict["model_call"]["finish_reason"] == "stop"
        assert span_dict.get("cost_usd") is not None
