"""
tests/test_google_embeddings.py — Tests for Google/Gemini embedding provider.

Tests cover:
    - Provider initialization and API key resolution
    - Batch embedding via mock HTTP (OpenAI-compatible endpoint)
    - Dimension configuration
    - Error handling (rate limit, auth failure)
    - Model detection helpers
    - HostProvider routing to Google embeddings
    - ProviderRouter configure_from_host with Google provider
"""

from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.llm.google_embeddings import (
    GOOGLE_BATCH_LIMIT,
    GOOGLE_EMBED_BASE_URL,
    GOOGLE_EMBED_MODELS,
    embed_google,
    get_google_api_key,
    is_google_embed_model,
)
from cortex.llm.host_provider import HostProvider, _is_google_embed_model
from cortex.llm.provider import EmbedResponse


# ---------------------------------------------------------------------------
# Fixtures & helpers
# ---------------------------------------------------------------------------

def _make_mock_embedding_response(texts: list[str], dim: int = 3072):
    """Create a mock OpenAI embeddings response."""
    resp = MagicMock()
    resp.data = []
    for i, _ in enumerate(texts):
        item = MagicMock()
        item.embedding = [0.1 * (i + 1)] * dim
        resp.data.append(item)
    resp.usage = MagicMock()
    resp.usage.prompt_tokens = len(texts) * 10
    return resp


def _make_mock_client(dim: int = 3072):
    """Create a mock AsyncOpenAI client for Google embedding endpoint."""
    client = MagicMock()
    client.embeddings = MagicMock()

    async def mock_create(**kwargs):
        texts = kwargs.get("input", [])
        return _make_mock_embedding_response(texts, dim)

    client.embeddings.create = mock_create
    return client


# ---------------------------------------------------------------------------
# Tests: model detection
# ---------------------------------------------------------------------------

class TestModelDetection:
    def test_known_models(self):
        assert is_google_embed_model("gemini-embedding-001")
        assert is_google_embed_model("gemini-embedding-2-preview")

    def test_prefix_match(self):
        assert is_google_embed_model("gemini-embedding-future-model")

    def test_non_google_models(self):
        assert not is_google_embed_model("voyage-code-3")
        assert not is_google_embed_model("text-embedding-3-small")
        assert not is_google_embed_model("gemini-pro")  # not embedding

    def test_host_provider_detection(self):
        assert _is_google_embed_model("gemini-embedding-001")
        assert _is_google_embed_model("gemini-embedding-2-preview")
        assert not _is_google_embed_model("voyage-code-3")


# ---------------------------------------------------------------------------
# Tests: API key resolution
# ---------------------------------------------------------------------------

class TestAPIKeyResolution:
    def test_from_config(self):
        key = get_google_api_key({"api_key": "config-key-123"})
        assert key == "config-key-123"

    def test_from_google_env(self, monkeypatch):
        monkeypatch.setenv("GOOGLE_API_KEY", "google-env-key")
        monkeypatch.delenv("GEMINI_API_KEY", raising=False)
        key = get_google_api_key()
        assert key == "google-env-key"

    def test_from_gemini_env(self, monkeypatch):
        monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
        monkeypatch.setenv("GEMINI_API_KEY", "gemini-env-key")
        key = get_google_api_key()
        assert key == "gemini-env-key"

    def test_config_takes_priority(self, monkeypatch):
        monkeypatch.setenv("GOOGLE_API_KEY", "env-key")
        key = get_google_api_key({"api_key": "config-key"})
        assert key == "config-key"

    def test_google_env_over_gemini(self, monkeypatch):
        monkeypatch.setenv("GOOGLE_API_KEY", "google-key")
        monkeypatch.setenv("GEMINI_API_KEY", "gemini-key")
        key = get_google_api_key()
        assert key == "google-key"

    def test_empty_without_keys(self, monkeypatch):
        monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
        monkeypatch.delenv("GEMINI_API_KEY", raising=False)
        key = get_google_api_key()
        assert key == ""


# ---------------------------------------------------------------------------
# Tests: embed_google()
# ---------------------------------------------------------------------------

class TestEmbedGoogle:
    @pytest.mark.asyncio
    async def test_basic_embedding(self):
        client = _make_mock_client(dim=3072)
        result = await embed_google(
            texts=["hello world"],
            model="gemini-embedding-001",
            api_key="test-key",
            _client=client,
        )
        assert isinstance(result, EmbedResponse)
        assert len(result.embeddings) == 1
        assert len(result.embeddings[0]) == 3072
        assert result.model == "gemini-embedding-001"
        assert result.tokens_in == 10

    @pytest.mark.asyncio
    async def test_multiple_texts(self):
        client = _make_mock_client(dim=1536)
        texts = ["text one", "text two", "text three"]
        result = await embed_google(
            texts=texts,
            model="gemini-embedding-2-preview",
            api_key="test-key",
            dim=1536,
            _client=client,
        )
        assert len(result.embeddings) == 3
        assert len(result.embeddings[0]) == 1536

    @pytest.mark.asyncio
    async def test_dimension_passed_to_api(self):
        """Verify dimensions kwarg is forwarded to the API call."""
        client = MagicMock()
        call_kwargs = {}

        async def capture_create(**kwargs):
            call_kwargs.update(kwargs)
            return _make_mock_embedding_response(kwargs["input"], 768)

        client.embeddings = MagicMock()
        client.embeddings.create = capture_create

        await embed_google(
            texts=["test"],
            model="gemini-embedding-2-preview",
            api_key="test-key",
            dim=768,
            _client=client,
        )
        assert call_kwargs.get("dimensions") == 768

    @pytest.mark.asyncio
    async def test_no_dimension_when_none(self):
        """When dim=None, dimensions kwarg should not be sent."""
        client = MagicMock()
        call_kwargs = {}

        async def capture_create(**kwargs):
            call_kwargs.update(kwargs)
            return _make_mock_embedding_response(kwargs["input"], 3072)

        client.embeddings = MagicMock()
        client.embeddings.create = capture_create

        await embed_google(
            texts=["test"],
            model="gemini-embedding-001",
            api_key="test-key",
            dim=None,
            _client=client,
        )
        assert "dimensions" not in call_kwargs

    @pytest.mark.asyncio
    async def test_batching(self):
        """Texts exceeding GOOGLE_BATCH_LIMIT should be split into batches."""
        batch_sizes = []
        client = MagicMock()

        async def track_create(**kwargs):
            batch_sizes.append(len(kwargs["input"]))
            return _make_mock_embedding_response(kwargs["input"], 3072)

        client.embeddings = MagicMock()
        client.embeddings.create = track_create

        # 150 texts → 2 batches (100 + 50)
        texts = [f"text {i}" for i in range(150)]

        with patch("cortex.llm.google_embeddings.asyncio.sleep", new_callable=AsyncMock):
            result = await embed_google(
                texts=texts,
                model="gemini-embedding-001",
                api_key="test-key",
                _client=client,
            )

        assert len(result.embeddings) == 150
        assert batch_sizes == [100, 50]

    @pytest.mark.asyncio
    async def test_batch_delay(self):
        """Verify delay between batches is applied."""
        sleep_calls = []
        client = _make_mock_client(dim=3072)

        async def mock_sleep(seconds):
            sleep_calls.append(seconds)

        texts = [f"text {i}" for i in range(150)]

        with patch("cortex.llm.google_embeddings.asyncio.sleep", side_effect=mock_sleep):
            await embed_google(
                texts=texts,
                model="gemini-embedding-001",
                api_key="test-key",
                _client=client,
            )

        # Should sleep once between batch 1 and batch 2
        assert len(sleep_calls) == 1
        assert sleep_calls[0] == 0.5

    @pytest.mark.asyncio
    async def test_cost_estimate_zero_for_free_tier(self):
        client = _make_mock_client(dim=3072)
        result = await embed_google(
            texts=["test"],
            model="gemini-embedding-001",
            api_key="test-key",
            _client=client,
        )
        assert result.cost_estimate == 0.0

    @pytest.mark.asyncio
    async def test_latency_tracked(self):
        client = _make_mock_client(dim=3072)
        result = await embed_google(
            texts=["test"],
            model="gemini-embedding-001",
            api_key="test-key",
            _client=client,
        )
        assert result.latency_ms >= 0

    @pytest.mark.asyncio
    async def test_empty_texts(self):
        client = _make_mock_client(dim=3072)
        result = await embed_google(
            texts=[],
            model="gemini-embedding-001",
            api_key="test-key",
            _client=client,
        )
        assert result.embeddings == []
        assert result.tokens_in == 0


# ---------------------------------------------------------------------------
# Tests: error handling
# ---------------------------------------------------------------------------

class TestErrorHandling:
    @pytest.mark.asyncio
    async def test_missing_api_key_raises(self, monkeypatch):
        monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
        monkeypatch.delenv("GEMINI_API_KEY", raising=False)

        with pytest.raises(RuntimeError, match="No Google API key"):
            from cortex.llm.google_embeddings import _get_google_embed_client
            _get_google_embed_client("")

    @pytest.mark.asyncio
    async def test_auth_failure(self):
        """Simulate 401 auth error from API."""
        client = MagicMock()

        async def raise_auth_error(**kwargs):
            raise Exception("Error code: 401 - Invalid API key")

        client.embeddings = MagicMock()
        client.embeddings.create = raise_auth_error

        with pytest.raises(Exception, match="401"):
            await embed_google(
                texts=["test"],
                model="gemini-embedding-001",
                api_key="bad-key",
                _client=client,
            )

    @pytest.mark.asyncio
    async def test_rate_limit_error(self):
        """Simulate 429 rate limit error from API."""
        client = MagicMock()

        async def raise_rate_limit(**kwargs):
            raise Exception("Error code: 429 - Rate limit exceeded")

        client.embeddings = MagicMock()
        client.embeddings.create = raise_rate_limit

        with pytest.raises(Exception, match="429"):
            await embed_google(
                texts=["test"],
                model="gemini-embedding-001",
                api_key="test-key",
                _client=client,
            )


# ---------------------------------------------------------------------------
# Tests: HostProvider integration
# ---------------------------------------------------------------------------

class TestHostProviderIntegration:
    @pytest.mark.asyncio
    async def test_host_provider_routes_to_google(self):
        """HostProvider.embed() should route gemini models to _embed_google."""
        provider = HostProvider(
            llm_config={"api_key": "llm-key", "model": "gpt-4.1-nano"},
            embedding_config={
                "api_key": "google-key",
                "model": "gemini-embedding-001",
                "dim": 3072,
                "provider": "google",
            },
        )

        mock_client = _make_mock_client(dim=3072)

        with patch(
            "cortex.llm.google_embeddings._get_google_embed_client",
            return_value=mock_client,
        ):
            result = await provider.embed(["test text"])

        assert isinstance(result, EmbedResponse)
        assert len(result.embeddings) == 1
        assert result.model == "gemini-embedding-001"

    @pytest.mark.asyncio
    async def test_host_provider_still_routes_voyage(self):
        """Voyage models should still route to Voyage, not Google."""
        provider = HostProvider(
            llm_config={"api_key": "llm-key", "model": "gpt-4.1-nano"},
            embedding_config={
                "api_key": "voyage-key",
                "model": "voyage-code-3",
                "dim": 1024,
            },
        )
        # Should NOT try to use Google — verify it calls _embed_voyage
        with patch.object(provider, "_embed_voyage", new_callable=AsyncMock) as mock_voyage:
            mock_voyage.return_value = EmbedResponse(
                embeddings=[[0.1] * 1024], model="voyage-code-3",
                tokens_in=10, tokens_out=0, cost_estimate=0.0, latency_ms=1.0,
            )
            result = await provider.embed(["test"], model="voyage-code-3")
            mock_voyage.assert_called_once()

    @pytest.mark.asyncio
    async def test_host_provider_explicit_model_override(self):
        """Passing model= to embed() should override default."""
        provider = HostProvider(
            llm_config={"api_key": "llm-key", "model": "gpt-4.1-nano"},
            embedding_config={
                "api_key": "google-key",
                "model": "voyage-code-3",  # default is voyage
                "dim": 1024,
            },
        )

        mock_client = _make_mock_client(dim=3072)

        with patch(
            "cortex.llm.google_embeddings._get_google_embed_client",
            return_value=mock_client,
        ):
            # Override with gemini model
            result = await provider.embed(["test"], model="gemini-embedding-001")

        assert result.model == "gemini-embedding-001"


# ---------------------------------------------------------------------------
# Tests: ProviderRouter configure_from_host with Google
# ---------------------------------------------------------------------------

class TestProviderRouterGoogle:
    def test_configure_from_host_with_google(self):
        """ProviderRouter should recognize 'google' as an embedding provider."""
        from cortex.llm.provider_router import ProviderRouter

        @dataclass
        class FakeConfig:
            provider_mode: str = "host"
            provider_llm_api_key: str = ""
            provider_llm_base_url: str = ""
            provider_llm_model: str = "gpt-4.1-nano"
            provider_llm_provider: str = ""
            provider_embedding_api_key: str = ""
            provider_embedding_base_url: str = ""
            provider_embedding_model: str = "voyage-code-3"
            provider_embedding_dim: int = 1024
            provider_embedding_provider: str = ""

        router = ProviderRouter(FakeConfig())

        configured = router.configure_from_host({
            "providers": {
                "openai": {"api_key": "oai-key"},
                "google": {"api_key": "google-key"},
            },
            "extraction_model": "gemini-2.5-flash",
            "embedding_model": "gemini-embedding-001",
            "embedding_dim": 3072,
        })

        # Google is now first in LLM preference, so it handles both LLM and embedding
        assert "google" in configured

    def test_google_preferred_over_openai_for_embeddings(self):
        """When both google and openai are available, google should be picked first for embeddings."""
        from cortex.llm.provider_router import ProviderRouter

        @dataclass
        class FakeConfig:
            provider_mode: str = "host"
            provider_llm_api_key: str = ""
            provider_llm_base_url: str = ""
            provider_llm_model: str = "gpt-4.1-nano"
            provider_llm_provider: str = ""
            provider_embedding_api_key: str = ""
            provider_embedding_base_url: str = ""
            provider_embedding_model: str = "voyage-code-3"
            provider_embedding_dim: int = 1024
            provider_embedding_provider: str = ""

        router = ProviderRouter(FakeConfig())

        configured = router.configure_from_host({
            "providers": {
                "openai": {"api_key": "oai-key"},
                "google": {"api_key": "google-key"},
            },
            "extraction_model": "gpt-4.1-nano",
            "embedding_model": "gemini-embedding-2-preview",
            "embedding_dim": 1536,
        })

        # Google should be in the configured list
        assert "google" in configured
        # The inner provider's embedding config should use google's key
        inner = router.provider
        assert inner._embedding_config["api_key"] == "google-key"
        assert inner._embedding_config["provider"] == "google"
