"""Tests of pydantic-ai actually connecting to OpenAI and Gemini models.

WARNING: running these tests will consume your OpenAI and Gemini credits.
"""

import os
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Callable

import httpx
import pytest
from pydantic import BaseModel

from pydantic_ai import Agent
from pydantic_ai.models import Model

pytestmark = [
    pytest.mark.skipif(os.getenv('PYDANTIC_AI_LIVE_TEST_DANGEROUS') != 'CHARGE-ME!', reason='live tests disabled'),
    pytest.mark.anyio,
]


def openai(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.openai import OpenAIModel

    return OpenAIModel('gpt-4o-mini', http_client=http_client)


def gemini(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.gemini import GeminiModel

    return GeminiModel('gemini-1.5-pro', http_client=http_client)


def vertexai(http_client: httpx.AsyncClient, tmp_path: Path) -> Model:
    from pydantic_ai.models.vertexai import VertexAIModel

    service_account_content = os.environ['GOOGLE_SERVICE_ACCOUNT_CONTENT']
    service_account_path = tmp_path / 'service_account.json'
    service_account_path.write_text(service_account_content)
    return VertexAIModel('gemini-1.5-flash', service_account_file=service_account_path, http_client=http_client)


def groq(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.groq import GroqModel

    return GroqModel('llama-3.3-70b-versatile', http_client=http_client)


def anthropic(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.anthropic import AnthropicModel

    return AnthropicModel('claude-3-5-sonnet-latest', http_client=http_client)


def ollama(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.openai import OpenAIModel

    return OpenAIModel('qwen2:0.5b', base_url='http://localhost:11434/v1/', http_client=http_client)


def mistral(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.mistral import MistralModel

    return MistralModel('mistral-small-latest', http_client=http_client)


def cohere(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
    from pydantic_ai.models.cohere import CohereModel

    return CohereModel('command-r7b-12-2024', http_client=http_client)


params = [
    pytest.param(openai, id='openai'),
    pytest.param(gemini, marks=pytest.mark.skip(reason='API seems very flaky'), id='gemini'),
    pytest.param(vertexai, id='vertexai'),
    pytest.param(groq, id='groq'),
    pytest.param(anthropic, id='anthropic'),
    pytest.param(ollama, id='ollama'),
    pytest.param(mistral, id='mistral'),
    pytest.param(cohere, id='cohere'),
]
GetModel = Callable[[httpx.AsyncClient, Path], Model]


@pytest.fixture
async def http_client(allow_model_requests: None) -> AsyncIterator[httpx.AsyncClient]:
    async with httpx.AsyncClient(timeout=30) as client:
        yield client


@pytest.mark.parametrize('get_model', params)
async def test_text(http_client: httpx.AsyncClient, tmp_path: Path, get_model: GetModel):
    agent = Agent(get_model(http_client, tmp_path), model_settings={'temperature': 0.0}, retries=2)
    result = await agent.run('What is the capital of France?')
    print('Text response:', result.data)
    assert 'paris' in result.data.lower()
    print('Text usage:', result.usage())
    usage = result.usage()
    assert usage.total_tokens is not None and usage.total_tokens > 0


stream_params = [p for p in params if p.id != 'cohere']


@pytest.mark.parametrize('get_model', stream_params)
async def test_stream(http_client: httpx.AsyncClient, tmp_path: Path, get_model: GetModel):
    agent = Agent(get_model(http_client, tmp_path), model_settings={'temperature': 0.0}, retries=2)
    async with agent.run_stream('What is the capital of France?') as result:
        data = await result.get_data()
    print('Stream response:', data)
    assert 'paris' in data.lower()
    print('Stream usage:', result.usage())
    usage = result.usage()
    if get_model.__name__ != 'ollama':
        assert usage.total_tokens is not None and usage.total_tokens > 0


class MyModel(BaseModel):
    city: str


structured_params = [p for p in params if p.id != 'ollama']


@pytest.mark.parametrize('get_model', structured_params)
async def test_structured(http_client: httpx.AsyncClient, tmp_path: Path, get_model: GetModel):
    agent = Agent(get_model(http_client, tmp_path), result_type=MyModel, model_settings={'temperature': 0.0}, retries=2)
    result = await agent.run('What is the capital of the UK?')
    print('Structured response:', result.data)
    assert result.data.city.lower() == 'london'
    print('Structured usage:', result.usage())
    usage = result.usage()
    assert usage.total_tokens is not None and usage.total_tokens > 0
