"""
Tests for Google embeddings implementation (Gemini API + Vertex AI).

These tests cover:
1. Initialization (Gemini API key, Vertex AI with ADC/service account)
2. Dimension detection via test embedding
3. Output dimensionality configuration
4. Encode (single text, multiple texts, batching, empty list, uninitialized)
5. Provider name and model name normalization
6. Factory function (create from env, validation errors)
"""

from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from hindsight_api.config import (
    ENV_EMBEDDINGS_GEMINI_API_KEY,
    ENV_EMBEDDINGS_PROVIDER,
    HindsightConfig,
)
from hindsight_api.engine.embeddings import GeminiEmbeddings, create_embeddings_from_env


def _make_mock_embedding(values: list[float]) -> MagicMock:
    emb = MagicMock()
    emb.values = values
    return emb


def _make_mock_embed_result(embeddings_data: list[list[float]]) -> MagicMock:
    result = MagicMock()
    result.embeddings = [_make_mock_embedding(v) for v in embeddings_data]
    return result


def _make_mock_genai(embed_result: Any = None) -> MagicMock:
    if embed_result is None:
        embed_result = _make_mock_embed_result([[0.1] * 768])
    mock_genai = MagicMock()
    mock_client = MagicMock()
    mock_client.models.embed_content = MagicMock(return_value=embed_result)
    mock_genai.Client = MagicMock(return_value=mock_client)
    return mock_genai


def _make_mock_google_module(mock_genai: MagicMock) -> MagicMock:
    mod = MagicMock()
    mod.genai = mock_genai
    mod.genai.types.EmbedContentConfig = MagicMock(side_effect=lambda **kw: MagicMock(**kw))
    return mod


def _patch_google_import(mock_genai: MagicMock):
    original_import = __import__

    def mock_import(name, *args, **kwargs):
        if name == "google":
            return _make_mock_google_module(mock_genai)
        if name == "google.genai":
            return mock_genai
        return original_import(name, *args, **kwargs)

    return patch("builtins.__import__", side_effect=mock_import)


class TestGeminiEmbeddings:
    """Unit tests for GeminiEmbeddings with mocked google.genai."""

    async def test_initialization_api_key_success(self):
        """Test successful Gemini API key initialization."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")

        with _patch_google_import(mock_genai):
            await emb.initialize()

        assert emb._client is not None
        assert emb.dimension == 768
        assert emb.provider_name == "google"
        assert emb._is_vertexai is False
        mock_genai.Client.return_value.models.embed_content.assert_called_once()

    async def test_initialization_vertexai_success(self):
        """Test successful Vertex AI initialization."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(
            model="gemini-embedding-001",
            vertexai_project_id="test-project",
            vertexai_region="us-central1",
        )

        with _patch_google_import(mock_genai):
            await emb.initialize()

        assert emb._client is not None
        assert emb.dimension == 768
        assert emb.provider_name == "google"
        assert emb._is_vertexai is True
        mock_genai.Client.assert_called_once_with(
            vertexai=True,
            project="test-project",
            location="us-central1",
        )

    async def test_initialization_missing_api_key(self):
        """Test that missing API key raises ValueError when no vertexai_project_id."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key=None)

        with _patch_google_import(mock_genai):
            with pytest.raises(ValueError, match="requires an API key"):
                await emb.initialize()

    async def test_initialization_vertexai_missing_project_id(self):
        """Test that Vertex AI mode requires project_id."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="gemini-embedding-001", vertexai_project_id="temp")
        emb.vertexai_project_id = None  # Simulate misconfiguration

        with _patch_google_import(mock_genai):
            with pytest.raises(ValueError, match="is required for Vertex AI"):
                await emb.initialize()

    async def test_initialization_idempotent(self):
        """Test that calling initialize() twice is a no-op."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")

        with _patch_google_import(mock_genai):
            await emb.initialize()
            first_client = emb._client
            await emb.initialize()
            assert emb._client is first_client

    async def test_dimension_detection_via_test_embedding(self):
        """Test that dimension is detected via a test embedding call."""
        test_embed = _make_mock_embed_result([[0.5] * 256])
        mock_genai = _make_mock_genai(embed_result=test_embed)
        emb = GeminiEmbeddings(model="some-new-model", api_key="test-key")

        with _patch_google_import(mock_genai):
            await emb.initialize()

        assert emb.dimension == 256

    async def test_output_dimensionality(self):
        """Test that output_dimensionality is passed via EmbedContentConfig."""
        test_embed = _make_mock_embed_result([[0.1] * 256])
        mock_genai = _make_mock_genai(embed_result=test_embed)
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", output_dimensionality=256)

        with _patch_google_import(mock_genai):
            await emb.initialize()

        assert emb.dimension == 256
        assert emb._embed_config is not None
        call_kwargs = mock_genai.Client.return_value.models.embed_content.call_args
        assert "config" in call_kwargs.kwargs

    async def test_no_output_dimensionality(self):
        """Test that no EmbedContentConfig is built when output_dimensionality is None."""
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", output_dimensionality=None)

        with _patch_google_import(mock_genai):
            await emb.initialize()

        assert emb._embed_config is None
        call_kwargs = mock_genai.Client.return_value.models.embed_content.call_args
        assert "config" not in call_kwargs.kwargs

    def test_auto_detect_vertexai(self):
        """Test that _is_vertexai is auto-detected from vertexai_project_id."""
        assert GeminiEmbeddings(model="m", api_key="k")._is_vertexai is False
        assert GeminiEmbeddings(model="m", vertexai_project_id="p")._is_vertexai is True

    def test_encode_single_text(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        mock_client = MagicMock()
        mock_client.models.embed_content = MagicMock(return_value=_make_mock_embed_result([[0.1, 0.2, 0.3]]))
        emb._client = mock_client
        emb._dimension = 3

        assert emb.encode(["hello"]) == [[0.1, 0.2, 0.3]]

    def test_encode_multiple_texts(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        mock_client = MagicMock()
        mock_client.models.embed_content = MagicMock(
            return_value=_make_mock_embed_result([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
        )
        emb._client = mock_client
        emb._dimension = 2

        result = emb.encode(["a", "b", "c"])
        assert len(result) == 3
        assert result[1] == [0.3, 0.4]

    def test_encode_batching(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", batch_size=2)
        mock_client = MagicMock()
        mock_client.models.embed_content = MagicMock(
            side_effect=[_make_mock_embed_result([[0.1], [0.2]]), _make_mock_embed_result([[0.3]])]
        )
        emb._client = mock_client
        emb._dimension = 1

        assert emb.encode(["a", "b", "c"]) == [[0.1], [0.2], [0.3]]
        assert mock_client.models.embed_content.call_count == 2

    def test_encode_passes_config(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        mock_client = MagicMock()
        mock_client.models.embed_content = MagicMock(return_value=_make_mock_embed_result([[0.1, 0.2]]))
        emb._client = mock_client
        emb._dimension = 2
        emb._embed_config = MagicMock()

        emb.encode(["hello"])
        assert mock_client.models.embed_content.call_args.kwargs["config"] is emb._embed_config

    def test_encode_empty_list(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        emb._client = MagicMock()
        emb._dimension = 768
        assert emb.encode([]) == []

    def test_encode_before_initialization(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        with pytest.raises(RuntimeError, match="not initialized"):
            emb.encode(["test"])

    def test_dimension_before_initialization(self):
        emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key")
        with pytest.raises(RuntimeError, match="not initialized"):
            _ = emb.dimension

    def test_provider_name_always_google(self):
        assert GeminiEmbeddings(model="m", api_key="k").provider_name == "google"
        assert GeminiEmbeddings(model="m", vertexai_project_id="p").provider_name == "google"

    def test_vertexai_strips_google_prefix(self):
        mock_genai = _make_mock_genai()
        emb = GeminiEmbeddings(model="google/gemini-embedding-001", vertexai_project_id="test-project")
        emb._init_vertexai(mock_genai)
        assert emb.model == "gemini-embedding-001"

    def test_default_region(self):
        emb = GeminiEmbeddings(model="m", vertexai_project_id="proj")
        assert emb.vertexai_region == "us-central1"

    def test_custom_region(self):
        emb = GeminiEmbeddings(model="m", vertexai_project_id="proj", vertexai_region="europe-west1")
        assert emb.vertexai_region == "europe-west1"


class TestGeminiEmbeddingsFactory:
    """Tests for create_embeddings_from_env() with 'google' provider."""

    def _make_config(self, **overrides) -> HindsightConfig:
        from dataclasses import fields

        defaults = {}
        for f in fields(HindsightConfig):
            if f.type == "str":
                defaults[f.name] = ""
            elif f.type == "str | None":
                defaults[f.name] = None
            elif f.type == "int":
                defaults[f.name] = 0
            elif f.type == "int | None":
                defaults[f.name] = None
            elif f.type == "float":
                defaults[f.name] = 0.0
            elif f.type == "float | None":
                defaults[f.name] = None
            elif f.type == "bool":
                defaults[f.name] = False
            elif f.type == "list | None":
                defaults[f.name] = None
            else:
                defaults[f.name] = None

        defaults["embeddings_provider"] = "google"
        defaults["embeddings_gemini_api_key"] = "test-key"
        defaults["embeddings_gemini_model"] = "gemini-embedding-001"
        defaults["embeddings_gemini_output_dimensionality"] = 768
        defaults["embeddings_vertexai_project_id"] = None
        defaults["embeddings_vertexai_region"] = None
        defaults["embeddings_vertexai_service_account_key"] = None

        defaults.update(overrides)
        return HindsightConfig(**defaults)

    def test_create_with_api_key(self):
        config = self._make_config()
        with patch("hindsight_api.config.get_config", return_value=config):
            emb = create_embeddings_from_env()
        assert isinstance(emb, GeminiEmbeddings)
        assert emb.provider_name == "google"
        assert emb.api_key == "test-key"
        assert emb._is_vertexai is False

    def test_create_with_vertexai(self):
        config = self._make_config(
            embeddings_gemini_api_key=None,
            embeddings_vertexai_project_id="my-project",
            embeddings_vertexai_region="us-east1",
        )
        with patch("hindsight_api.config.get_config", return_value=config):
            emb = create_embeddings_from_env()
        assert isinstance(emb, GeminiEmbeddings)
        assert emb._is_vertexai is True
        assert emb.api_key is None
        assert emb.vertexai_project_id == "my-project"

    def test_create_missing_all_credentials(self):
        config = self._make_config(embeddings_gemini_api_key=None, embeddings_vertexai_project_id=None)
        with patch("hindsight_api.config.get_config", return_value=config):
            with pytest.raises(ValueError, match="is required"):
                create_embeddings_from_env()

    def test_vertexai_takes_priority(self):
        config = self._make_config(embeddings_gemini_api_key="key", embeddings_vertexai_project_id="proj")
        with patch("hindsight_api.config.get_config", return_value=config):
            emb = create_embeddings_from_env()
        assert emb._is_vertexai is True
        assert emb.api_key is None

    def test_create_with_custom_dimensionality(self):
        config = self._make_config(embeddings_gemini_output_dimensionality=256)
        with patch("hindsight_api.config.get_config", return_value=config):
            emb = create_embeddings_from_env()
        assert emb.output_dimensionality == 256