# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

"""Tests for Google GenAI plugin."""

import os
from unittest.mock import MagicMock, patch

import pytest

from genkit.core.registry import ActionKind
from genkit.plugins.google_genai import (
    EmbeddingTaskType,
    GeminiConfigSchema,
    GeminiEmbeddingModels,
    GoogleAI,
    GoogleAIGeminiVersion,
    VertexAI,
    VertexAIGeminiVersion,
    VertexEmbeddingModels,
)
from genkit.plugins.google_genai.google import (
    GOOGLEAI_PLUGIN_NAME,
    VERTEXAI_PLUGIN_NAME,
    GenaiModels,
    googleai_name,
    vertexai_name,
)


def test_googleai_name() -> None:
    """Test googleai_name helper function."""
    assert googleai_name('gemini-2.0-flash') == 'googleai/gemini-2.0-flash'
    assert googleai_name('text-embedding-004') == 'googleai/text-embedding-004'


def test_vertexai_name() -> None:
    """Test vertexai_name helper function."""
    assert vertexai_name('gemini-2.0-flash') == 'vertexai/gemini-2.0-flash'
    assert vertexai_name('imagen-3.0-generate-001') == 'vertexai/imagen-3.0-generate-001'


def test_plugin_names() -> None:
    """Test plugin name constants."""
    assert GOOGLEAI_PLUGIN_NAME == 'googleai'
    assert VERTEXAI_PLUGIN_NAME == 'vertexai'


def test_googleai_initialization_with_api_key() -> None:
    """Test GoogleAI plugin initializes with API key parameter."""
    with patch('genkit.plugins.google_genai.google.genai.client.Client'):
        plugin = GoogleAI(api_key='test-key')
        assert plugin.name == 'googleai'
        assert plugin._vertexai is False


def test_googleai_initialization_from_env() -> None:
    """Test GoogleAI plugin reads API key from environment."""
    with patch.dict(os.environ, {'GEMINI_API_KEY': 'env-key'}):
        with patch('genkit.plugins.google_genai.google.genai.client.Client'):
            plugin = GoogleAI()
            assert plugin.name == 'googleai'


def test_googleai_initialization_without_api_key() -> None:
    """Test GoogleAI plugin raises error without API key."""
    with patch.dict(os.environ, {}, clear=True):
        with pytest.raises(ValueError) as exc_info:
            GoogleAI()
        assert 'GEMINI_API_KEY' in str(exc_info.value)


def test_vertexai_initialization() -> None:
    """Test VertexAI plugin initializes correctly."""
    with patch('genkit.plugins.google_genai.google.genai.client.Client'):
        plugin = VertexAI(project='test-project', location='us-central1')
        assert plugin.name == 'vertexai'
        assert plugin._vertexai is True


def test_vertexai_initialization_from_env() -> None:
    """Test VertexAI plugin reads project from environment."""
    with patch.dict(os.environ, {'GCLOUD_PROJECT': 'env-project'}):
        with patch('genkit.plugins.google_genai.google.genai.client.Client'):
            plugin = VertexAI()
            assert plugin.name == 'vertexai'


def test_genai_models_container() -> None:
    """Test GenaiModels container initialization."""
    models = GenaiModels()
    assert models.gemini == []
    assert models.imagen == []
    assert models.embedders == []
    assert models.veo == []


@patch('genkit.plugins.google_genai.google.genai.client.Client')
@patch('genkit.plugins.google_genai.google._list_genai_models')
@pytest.mark.asyncio
async def test_googleai_resolve_model(mock_list_models: MagicMock, mock_client: MagicMock) -> None:
    """Test GoogleAI plugin resolves model actions."""
    mock_list_models.return_value = GenaiModels()

    plugin = GoogleAI(api_key='test-key')
    action = await plugin.resolve(ActionKind.MODEL, 'googleai/gemini-2.0-flash')

    assert action is not None
    assert action.kind == ActionKind.MODEL
    assert action.name == 'googleai/gemini-2.0-flash'


@patch('genkit.plugins.google_genai.google.genai.client.Client')
@patch('genkit.plugins.google_genai.google._list_genai_models')
@pytest.mark.asyncio
async def test_googleai_resolve_embedder(mock_list_models: MagicMock, mock_client: MagicMock) -> None:
    """Test GoogleAI plugin resolves embedder actions."""
    mock_list_models.return_value = GenaiModels()

    plugin = GoogleAI(api_key='test-key')
    action = await plugin.resolve(ActionKind.EMBEDDER, 'googleai/text-embedding-004')

    assert action is not None
    assert action.kind == ActionKind.EMBEDDER
    assert action.name == 'googleai/text-embedding-004'


@patch('genkit.plugins.google_genai.google.genai.client.Client')
@patch('genkit.plugins.google_genai.google._list_genai_models')
@pytest.mark.asyncio
async def test_googleai_resolve_non_model_returns_none(mock_list_models: MagicMock, mock_client: MagicMock) -> None:
    """Test GoogleAI plugin returns None for unsupported action kinds."""
    mock_list_models.return_value = GenaiModels()

    plugin = GoogleAI(api_key='test-key')
    action = await plugin.resolve(ActionKind.PROMPT, 'some-prompt')
    assert action is None


@patch('genkit.plugins.google_genai.google.genai.client.Client')
@patch('genkit.plugins.google_genai.google._list_genai_models')
@pytest.mark.asyncio
async def test_vertexai_resolve_model(mock_list_models: MagicMock, mock_client: MagicMock) -> None:
    """Test VertexAI plugin resolves model actions."""
    mock_list_models.return_value = GenaiModels()

    plugin = VertexAI(project='test-project')
    action = await plugin.resolve(ActionKind.MODEL, 'vertexai/gemini-2.0-flash')

    assert action is not None
    assert action.kind == ActionKind.MODEL
    assert action.name == 'vertexai/gemini-2.0-flash'


def test_embedding_task_types() -> None:
    """Test EmbeddingTaskType enum values."""
    assert EmbeddingTaskType.RETRIEVAL_QUERY is not None
    assert EmbeddingTaskType.RETRIEVAL_DOCUMENT is not None
    assert EmbeddingTaskType.SEMANTIC_SIMILARITY is not None
    assert EmbeddingTaskType.CLASSIFICATION is not None
    assert EmbeddingTaskType.CLUSTERING is not None


def test_gemini_embedding_models_enum() -> None:
    """Test GeminiEmbeddingModels enum has values."""
    # Check that the enum has at least one value
    assert len(list(GeminiEmbeddingModels)) > 0


def test_vertex_embedding_models_enum() -> None:
    """Test VertexEmbeddingModels enum has values."""
    # Check that the enum has at least one value
    assert len(list(VertexEmbeddingModels)) > 0


def test_googleai_gemini_version_enum() -> None:
    """Test GoogleAIGeminiVersion enum has values."""
    # Check that the enum has at least one value
    assert len(list(GoogleAIGeminiVersion)) > 0


def test_vertexai_gemini_version_enum() -> None:
    """Test VertexAIGeminiVersion enum has values."""
    # Check that the enum has at least one value
    assert len(list(VertexAIGeminiVersion)) > 0


def test_gemini_config_schema() -> None:
    """Test GeminiConfigSchema can be instantiated."""
    config = GeminiConfigSchema(temperature=0.7, max_output_tokens=1000)
    assert config.temperature == 0.7
    assert config.max_output_tokens == 1000


def test_gemini_config_schema_defaults() -> None:
    """Test GeminiConfigSchema has proper defaults."""
    config = GeminiConfigSchema()
    # All fields should be optional with None defaults
    assert config.temperature is None
    assert config.max_output_tokens is None
