"""Tests for OpenAI auto-instrumentation.

Uses a mock OpenAI client to verify span capture without real API calls.
"""
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.tracing.context import trace


@dataclass
class FakeUsage:
    prompt_tokens: int = 100
    completion_tokens: int = 50
    total_tokens: int = 150


@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 TestOpenAIInstrumentation:
    @pytest.fixture
    def setup_openai(self):
        """Create a mock openai module structure for instrumentation."""
        # Build fake module structure matching openai SDK
        openai_mod = MagicMock()
        completions_cls = MagicMock()

        original_create = MagicMock(return_value=FakeResponse())
        completions_cls.create = original_create

        openai_mod.resources.chat.completions.Completions = completions_cls
        openai_mod.resources.chat.completions.AsyncCompletions = MagicMock()

        return openai_mod, completions_cls, original_create

    def test_instrument_patches_create(self, setup_openai, mock_fibinaci: MockFibinaci):
        openai_mod, completions_cls, original_create = setup_openai

        with patch.dict("sys.modules", {"openai": openai_mod, "openai.resources": openai_mod.resources,
                                         "openai.resources.chat": openai_mod.resources.chat,
                                         "openai.resources.chat.completions": openai_mod.resources.chat.completions}):
            from fibinaci.instrumentation.openai import OpenAIInstrumentor
            inst = OpenAIInstrumentor()
            inst.instrument()

            # The create method should now be wrapped
            assert completions_cls.create is not original_create

            inst.uninstrument()

    def test_span_captured_on_success(self, mock_fibinaci: MockFibinaci):
        """Verify a model_call span is pushed to the buffer on successful call."""
        from fibinaci.instrumentation._wrapper import _extract_span, ProviderConfig

        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,
        )
        response = FakeResponse(model="gpt-4o")
        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": "gpt-4o", "temperature": 0.7}, response, started, ended)

        assert span_dict["span_kind"] == "model_call"
        assert span_dict["provider"] == "openai"
        assert span_dict["tokens_input"] == 100
        assert span_dict["tokens_output"] == 50
        assert span_dict["status"] == "ok"
        assert span_dict["model_call"]["model_name"] == "gpt-4o"
        assert span_dict["duration_ms"] == 1000.0
        # Auto-cost
        assert span_dict.get("cost_usd") is not None
