"""
Unit tests for pydantic-ai integration
"""

import os
import pytest
from unittest.mock import patch, MagicMock

from latch.integrations.pydantic_ai import get_model, LatchModel, LatchModelSettings


@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set"
)
def test_get_model_openai_gpt():
    """Test get_model returns OpenAI model for GPT names"""
    models = [
        "gpt-4o-mini",
        "gpt-4o",
        "gpt-4-turbo",
        "gpt-4.1",
        "gpt-4.1-mini",
        "gpt-4.1-nano",
        "gpt-5",
        "gpt-5-mini",
        "gpt-5-nano",
    ]
    for model_name in models:
        model = get_model(model_name)
        assert "OpenAI" in type(model).__name__, f"Failed for {model_name}"


@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set"
)
def test_get_model_openai_reasoning():
    """Test get_model returns OpenAI model for o-series reasoning models"""
    models = [
        "o1",
        "o1-mini",
        "o3",
        "o3-mini",
        "o4-mini",
    ]
    for model_name in models:
        model = get_model(model_name)
        assert "OpenAI" in type(model).__name__, f"Failed for {model_name}"


@pytest.mark.skipif(
    not os.environ.get("ANTHROPIC_API_KEY"),
    reason="ANTHROPIC_API_KEY not set"
)
def test_get_model_anthropic():
    """Test get_model returns Anthropic model for Claude names"""
    models = [
        "claude-3-sonnet",
        "claude-3-opus",
        "claude-3-haiku",
        "claude-sonnet-4-5-20250514",
        "claude-opus-4-5-20250514",
        "claude-haiku-4-5-20250514",
        "claude-3-5-sonnet-20241022",
    ]
    for model_name in models:
        model = get_model(model_name)
        assert "Anthropic" in type(model).__name__, f"Failed for {model_name}"


@pytest.mark.skipif(
    not os.environ.get("GEMINI_API_KEY"),
    reason="GEMINI_API_KEY not set"
)
def test_get_model_gemini():
    """Test get_model returns Gemini model"""
    model = get_model("gemini-pro")
    assert "Gemini" in type(model).__name__

    model = get_model("gemini-1.5-flash")
    assert "Gemini" in type(model).__name__


def test_get_model_unknown():
    """Test get_model raises for unknown model"""
    with pytest.raises(ValueError, match="Unknown model"):
        get_model("unknown-model-xyz")


@pytest.mark.skipif(
    not os.environ.get("OPENAI_API_KEY") or not os.environ.get("ANTHROPIC_API_KEY"),
    reason="OPENAI_API_KEY and ANTHROPIC_API_KEY required"
)
def test_get_model_case_insensitive():
    """Test get_model is case insensitive"""
    model = get_model("GPT-4o-Mini")
    assert "OpenAI" in type(model).__name__

    model = get_model("CLAUDE-3-SONNET")
    assert "Anthropic" in type(model).__name__


def test_latch_model_settings():
    """Test LatchModelSettings TypedDict"""
    settings: LatchModelSettings = {
        "task": "qa",
        "budget": 0.10,
        "max_latency": 3000,
        "latch_metadata": {"user": "test"},
    }
    assert settings["task"] == "qa"
    assert settings["budget"] == 0.10


def test_latch_model_creation():
    """Test LatchModel creation"""
    model = LatchModel(
        api_key="test_key",
        base_url="http://localhost:8081",
        task="code-gen",
        budget=0.50,
        max_latency=5000,
    )

    assert model.api_key == "test_key"
    assert model.task == "code-gen"
    assert model.budget == 0.50
    assert model.max_latency == 5000
    assert model.model_name == "latch:code-gen"
    assert model.system == "latch"


def test_latch_model_defaults():
    """Test LatchModel default values"""
    model = LatchModel(api_key="test_key")

    assert model.base_url == "http://localhost:8081"
    assert model.task == "general"
    assert model.budget is None
    assert model.max_latency is None
    assert model.timeout == 60.0


def test_latch_model_convert_messages():
    """Test message conversion to OpenAI format"""
    model = LatchModel(api_key="test_key")

    class MockPart:
        part_kind = "user-prompt"
        content = "Hello world"

    class MockMessage:
        parts = [MockPart()]

    messages = model._convert_messages([MockMessage()])

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == "Hello world"


def test_latch_model_convert_messages_system():
    """Test system prompt conversion"""
    model = LatchModel(api_key="test_key")

    class MockPart:
        part_kind = "system-prompt"
        content = "You are helpful"

    class MockMessage:
        parts = [MockPart()]

    messages = model._convert_messages([MockMessage()])

    assert len(messages) == 1
    assert messages[0]["role"] == "system"


def test_latch_model_convert_messages_assistant():
    """Test assistant message conversion"""
    model = LatchModel(api_key="test_key")

    class MockPart:
        part_kind = "text"
        content = "I am the assistant"

    class MockMessage:
        parts = [MockPart()]

    messages = model._convert_messages([MockMessage()])

    assert len(messages) == 1
    assert messages[0]["role"] == "assistant"


def test_latch_model_get_settings():
    """Test settings extraction from model_settings dict"""
    model = LatchModel(api_key="test_key", task="default", budget=0.01)

    settings = {
        "task": "override",
        "budget": 0.50,
        "max_latency": 2000,
        "latch_metadata": {"key": "value"},
    }

    task, budget, max_latency, metadata = model._get_settings(settings)

    assert task == "override"
    assert budget == 0.50
    assert max_latency == 2000
    assert metadata == {"key": "value"}


def test_latch_model_get_settings_defaults():
    """Test settings falls back to model defaults"""
    model = LatchModel(api_key="test_key", task="default", budget=0.01)

    task, budget, max_latency, metadata = model._get_settings(None)

    assert task == "default"
    assert budget == 0.01
    assert max_latency is None
    assert metadata is None
