"""Unit tests for Gemini client embed()."""

from __future__ import annotations

import httpx
import pytest
import respx

from llm_rotator._types import EmbeddingResponse
from llm_rotator.clients.gemini import GeminiClient
from llm_rotator.exceptions import KeyDeadError, ServerError

GEMINI_EMBEDDING_RESPONSE = {
    "embeddings": [
        {"values": [0.01, 0.02, 0.03]},
        {"values": [0.04, 0.05, 0.06]},
    ]
}


class TestGeminiEmbed:
    @respx.mock
    async def test_successful_embed(self):
        respx.post(
            "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:batchEmbedContents"
        ).mock(return_value=httpx.Response(200, json=GEMINI_EMBEDDING_RESPONSE))

        client = GeminiClient()
        result = await client.embed(
            input=["hello", "world"],
            model="text-embedding-004",
            api_key="AIzaSyAB1234",
        )

        assert isinstance(result, EmbeddingResponse)
        assert len(result.embeddings) == 2
        assert result.embeddings[0] == [0.01, 0.02, 0.03]
        assert result.embeddings[1] == [0.04, 0.05, 0.06]
        assert result.model == "text-embedding-004"
        assert result.provider == "gemini"

    @respx.mock
    async def test_embed_request_format(self):
        """Verify Gemini batchEmbedContents request body format."""
        captured_body = {}

        def capture_request(request):
            import json

            captured_body.update(json.loads(request.content))
            return httpx.Response(200, json=GEMINI_EMBEDDING_RESPONSE)

        respx.post(
            "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:batchEmbedContents"
        ).mock(side_effect=capture_request)

        client = GeminiClient()
        await client.embed(
            input=["text1", "text2"],
            model="text-embedding-004",
            api_key="AIzaSyAB1234",
        )

        assert "requests" in captured_body
        assert len(captured_body["requests"]) == 2
        assert captured_body["requests"][0]["model"] == "models/text-embedding-004"
        assert captured_body["requests"][0]["content"] == {"parts": [{"text": "text1"}]}

    @respx.mock
    async def test_embed_custom_base_url(self):
        respx.post(
            "https://custom.gemini.api/v1beta/models/text-embedding-004:batchEmbedContents"
        ).mock(return_value=httpx.Response(200, json=GEMINI_EMBEDDING_RESPONSE))

        client = GeminiClient()
        result = await client.embed(
            input=["test"],
            model="text-embedding-004",
            api_key="AIzaSyAB1234",
            base_url="https://custom.gemini.api",
        )

        assert isinstance(result, EmbeddingResponse)

    @respx.mock
    async def test_embed_auth_error(self):
        respx.post(
            "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:batchEmbedContents"
        ).mock(return_value=httpx.Response(401, json={"error": {"message": "Bad key"}}))

        client = GeminiClient()
        with pytest.raises(KeyDeadError):
            await client.embed(
                input=["test"],
                model="text-embedding-004",
                api_key="bad-key-1",
            )

    @respx.mock
    async def test_embed_server_error(self):
        respx.post(
            "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:batchEmbedContents"
        ).mock(return_value=httpx.Response(500, json={"error": {"message": "Internal"}}))

        client = GeminiClient()
        with pytest.raises(ServerError):
            await client.embed(
                input=["test"],
                model="text-embedding-004",
                api_key="AIzaSyAB1234",
            )
