"""Tests for Vercel AI SDK 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


@dataclass
class FakeUsage:
    promptTokens: int = 75
    completionTokens: int = 35


@dataclass
class FakeVercelResponse:
    usage: FakeUsage = None
    finishReason: str = "stop"
    model: str = "gpt-4o"

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


class TestVercelAIInstrumentation:
    @pytest.fixture
    def setup_vercel(self):
        mod = MagicMock()
        original_generate = MagicMock(return_value=FakeVercelResponse())
        original_agenerate = MagicMock(return_value=FakeVercelResponse())
        mod.generate_text = original_generate
        mod.agenerate_text = original_agenerate
        return mod, original_generate

    def test_instrument_patches_generate_text(self, setup_vercel, mock_fibinaci: MockFibinaci):
        mod, original_generate = setup_vercel
        with patch.dict("sys.modules", {"ai": mod}):
            from fibinaci.instrumentation.vercel_ai import VercelAIInstrumentor
            inst = VercelAIInstrumentor()
            inst.instrument()
            assert mod.generate_text is not original_generate
            inst.uninstrument()

    def test_extract_vercel_usage_object(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.vercel_ai import _extract_vercel_usage

        response = FakeVercelResponse()
        inp, out, finish = _extract_vercel_usage(response)
        assert inp == 75
        assert out == 35
        assert finish == "stop"

    def test_extract_vercel_usage_dict(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.vercel_ai import _extract_vercel_usage

        response = {
            "usage": {"promptTokens": 100, "completionTokens": 50},
            "finishReason": "length",
        }
        inp, out, finish = _extract_vercel_usage(response)
        assert inp == 100
        assert out == 50
        assert finish == "length"

    def test_extract_vercel_usage_none(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.vercel_ai import _extract_vercel_usage

        @dataclass
        class EmptyResponse:
            usage: Any = None
            finishReason: Any = None

        inp, out, finish = _extract_vercel_usage(EmptyResponse())
        assert inp is None
        assert out is None
        assert finish is None
