"""Tests for Cohere 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:
    input_tokens: int = 70
    output_tokens: int = 25


@dataclass
class FakeResponse:
    model: str = "command-r-plus"
    usage: FakeUsage = None
    finish_reason: str = "COMPLETE"

    def __post_init__(self):
        if self.usage is None:
            self.usage = FakeUsage()


class TestCohereInstrumentation:
    @pytest.fixture
    def setup_cohere(self):
        mod = MagicMock()
        client_cls = MagicMock()
        async_cls = MagicMock()
        original_chat = MagicMock(return_value=FakeResponse())
        original_generate = MagicMock(return_value=FakeResponse())
        client_cls.chat = original_chat
        client_cls.generate = original_generate
        mod.Client = client_cls
        mod.AsyncClient = async_cls
        return mod, client_cls, original_chat

    def test_instrument_patches_chat(self, setup_cohere, mock_fibinaci: MockFibinaci):
        mod, client_cls, original_chat = setup_cohere
        with patch.dict("sys.modules", {"cohere": mod}):
            from fibinaci.instrumentation.cohere import CohereInstrumentor
            inst = CohereInstrumentor()
            inst.instrument()
            assert client_cls.chat is not original_chat
            inst.uninstrument()

    def test_span_captured_on_success(self, mock_fibinaci: MockFibinaci):
        config = ProviderConfig(
            name="cohere",
            input_tokens_field="input_tokens",
            output_tokens_field="output_tokens",
            finish_reason_getter=lambda r: getattr(r, "finish_reason", None),
        )
        response = FakeResponse()
        started = datetime(2025, 3, 29, 12, 0, 0, tzinfo=timezone.utc)
        ended = datetime(2025, 3, 29, 12, 0, 1, tzinfo=timezone.utc)

        span_dict = _extract_span(config, {"model": "command-r-plus"}, response, started, ended)

        assert span_dict["span_kind"] == "model_call"
        assert span_dict["provider"] == "cohere"
        assert span_dict["tokens_input"] == 70
        assert span_dict["tokens_output"] == 25
        assert span_dict["status"] == "ok"
        assert span_dict["model_call"]["finish_reason"] == "COMPLETE"

