"""Tests for Google Generative AI 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 FakeUsageMetadata:
    prompt_token_count: int = 60
    candidates_token_count: int = 30


@dataclass
class FakeGenResponse:
    usage_metadata: FakeUsageMetadata = None

    def __post_init__(self):
        if self.usage_metadata is None:
            self.usage_metadata = FakeUsageMetadata()


class TestGoogleGenAIInstrumentation:
    @pytest.fixture
    def setup_google(self):
        genai_mod = MagicMock()
        model_cls = MagicMock()
        original_generate = MagicMock(return_value=FakeGenResponse())
        original_agenerate = MagicMock(return_value=FakeGenResponse())
        model_cls.generate_content = original_generate
        model_cls.generate_content_async = original_agenerate
        genai_mod.GenerativeModel = model_cls

        google_mod = MagicMock()
        google_mod.generativeai = genai_mod
        return google_mod, genai_mod, model_cls, original_generate

    def test_instrument_patches_generate_content(self, setup_google, mock_fibinaci: MockFibinaci):
        google_mod, genai_mod, model_cls, original_generate = setup_google
        with patch.dict("sys.modules", {
            "google": google_mod,
            "google.generativeai": genai_mod,
        }):
            from fibinaci.instrumentation.google_genai import GoogleGenAIInstrumentor
            inst = GoogleGenAIInstrumentor()
            inst.instrument()
            assert model_cls.generate_content is not original_generate
            inst.uninstrument()

    def test_extract_google_usage(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.google_genai import _extract_google_usage

        response = FakeGenResponse()
        inp, out = _extract_google_usage(response)
        assert inp == 60
        assert out == 30

    def test_extract_google_usage_none(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.google_genai import _extract_google_usage

        @dataclass
        class EmptyResponse:
            usage_metadata: Any = None

        inp, out = _extract_google_usage(EmptyResponse())
        assert inp is None
        assert out is None
