# Copyright 2025 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

"""Unit-Tests for GoogleAI & VertexAI plugin."""

import asyncio
import os
import sys  # noqa
import unittest
from dataclasses import dataclass
from unittest.mock import ANY, MagicMock, patch

import pytest
from google import genai
from google.auth.credentials import Credentials
from google.genai.types import HttpOptions

from genkit import (
    ActionKind,
    Genkit,
    Message,
    ModelInfo,
    ModelRequest,
    Part,
    Role,
    TextPart,
)
from genkit.plugin_api import GENKIT_CLIENT_HEADER
from genkit.plugins.google_genai import GoogleAI, VertexAI
from genkit.plugins.google_genai.google import _inject_attribution_headers, googleai_name, vertexai_name
from genkit.plugins.google_genai.models.gemini import (
    DEFAULT_SUPPORTS_MODEL,
    SUPPORTED_MODELS,
    GeminiConfigSchema,
    GeminiModel,
)
from genkit.plugins.google_genai.models.imagen import (
    DEFAULT_IMAGE_SUPPORT,
    SUPPORTED_MODELS as IMAGE_SUPPORTED_MODELS,
)


async def _get_runtime_client(plugin: GoogleAI | VertexAI) -> object:
    return plugin._runtime_client()


@pytest.fixture
@patch('google.genai.client.Client')
def googleai_plugin_instance(client: MagicMock) -> GoogleAI:
    """GoogleAI fixture."""
    api_key = 'test_api_key'
    return GoogleAI(api_key=api_key)


class TestGoogleAIInit(unittest.TestCase):
    """Test cases for __init__ plugin."""

    @patch('google.genai.client.Client')
    def test_init_with_api_key(self, mock_genai_client: MagicMock) -> None:
        """Test using api_key parameter."""
        api_key = 'test_api_key'
        plugin = GoogleAI(api_key=api_key)
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=False,
            api_key=api_key,
            credentials=None,
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, GoogleAI)
        self.assertFalse(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)

    @patch('google.genai.client.Client')
    @patch.dict(os.environ, {'GEMINI_API_KEY': 'env_api_key'})
    def test_init_from_env_var(self, mock_genai_client: MagicMock) -> None:
        """Test using env var for api_key."""
        plugin = GoogleAI()
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=False,
            api_key='env_api_key',
            credentials=None,
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, GoogleAI)
        self.assertFalse(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)

    @patch('google.genai.client.Client')
    def test_init_with_credentials(self, mock_genai_client: MagicMock) -> None:
        """Test using credentials parameter."""
        mock_credentials = MagicMock(spec=Credentials)
        plugin = GoogleAI(credentials=mock_credentials)
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=False,
            api_key=ANY,
            credentials=mock_credentials,
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, GoogleAI)
        self.assertFalse(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)

    def test_init_raises_value_error_no_api_key(self) -> None:
        """Test using credentials parameter."""
        with (
            patch.dict(os.environ, {'GEMINI_API_KEY': ''}, clear=True),
            self.assertRaisesRegex(
                ValueError,
                r'GEMINI_API_KEY environment variable not set',
            ),
        ):
            GoogleAI()


@patch('google.genai.client.Client')
@pytest.mark.asyncio
async def test_googleai_initialize(mock_client_cls: MagicMock) -> None:
    """Unit tests for GoogleAI.init method."""
    mock_client = mock_client_cls.return_value

    m1 = MagicMock()
    m1.name = 'models/gemini-pro'
    m1.supported_actions = ['generateContent']
    m1.description = ' Gemini Pro '

    m2 = MagicMock()
    m2.name = 'models/gemini-embedding-001'
    m2.supported_actions = ['embedContent']
    m2.description = ' Embedding '

    mock_client.models.list.return_value = [m1, m2]

    api_key = 'test_api_key'
    plugin = GoogleAI(api_key=api_key)
    plugin._runtime_client = lambda: mock_client

    await plugin.init()
    result = await plugin.list_actions()

    # init returns known models and embedders
    assert len(result) > 0, 'Should initialize with known models and embedders'
    assert all(action.action_type is not None for action in result), 'All actions should have an action_type'
    assert all(hasattr(action, 'name') for action in result), 'All actions should have a name'
    assert all(action.name.startswith('googleai/') for action in result), (
        "All actions should be namespaced with 'googleai/'"
    )

    # Verify we have both models and embedders
    model_actions = [a for a in result if a.action_type == ActionKind.MODEL]
    embedder_actions = [a for a in result if a.action_type == ActionKind.EMBEDDER]
    assert len(model_actions) > 0, 'Should have at least one model'
    assert len(embedder_actions) > 0, 'Should have at least one embedder'


@patch('genkit.plugins.google_genai.GoogleAI._resolve_model')
@pytest.mark.asyncio
async def test_googleai_resolve_action_model(
    mock_resolve_action: MagicMock, googleai_plugin_instance: GoogleAI
) -> None:
    """Test resolve action for model."""
    plugin = googleai_plugin_instance

    await plugin.resolve(action_type=ActionKind.MODEL, name='lazaro-model')
    mock_resolve_action.assert_called_once_with('lazaro-model')


@patch('genkit.plugins.google_genai.GoogleAI._resolve_embedder')
@pytest.mark.asyncio
async def test_googleai_resolve_action_embedder(
    mock_resolve_action: MagicMock, googleai_plugin_instance: GoogleAI
) -> None:
    """Test resolve action for embedder."""
    plugin = googleai_plugin_instance

    await plugin.resolve(action_type=ActionKind.EMBEDDER, name='lazaro-model')
    mock_resolve_action.assert_called_once_with('lazaro-model')


@patch('genkit.plugins.google_genai.models.gemini.google_model_info')
@pytest.mark.parametrize(
    'model_name, expected_model_name, key',
    [
        (
            'gemini-pro-deluxe-max',
            'googleai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
        ),
        (
            'googleai/gemini-pro-deluxe-max',
            'googleai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
        ),
    ],
)
def test_googleai__resolve_model(
    mock_google_model_info: MagicMock,
    model_name: str,
    expected_model_name: str,
    key: str,
    googleai_plugin_instance: GoogleAI,
) -> None:
    """Tests for GoogleAI._resolve_model method."""
    plugin = googleai_plugin_instance

    mock_google_model_info.return_value = ModelInfo(
        label=f'Google AI - {model_name}',
        supports=DEFAULT_SUPPORTS_MODEL,
    )

    action = plugin._resolve_model(name=expected_model_name)

    assert action is not None
    assert action.kind == ActionKind.MODEL
    assert action.name == expected_model_name
    assert key in SUPPORTED_MODELS


@pytest.mark.parametrize(
    'model_name, expected_model_name, clean_name',
    [
        ('gemini-pro-deluxe-max', 'googleai/gemini-pro-deluxe-max', 'gemini-pro-deluxe-max'),
        ('googleai/gemini-pro-deluxe-max', 'googleai/gemini-pro-deluxe-max', 'gemini-pro-deluxe-max'),
    ],
)
def test_googleai__resolve_embedder(
    model_name: str,
    expected_model_name: str,
    clean_name: str,
    googleai_plugin_instance: GoogleAI,
) -> None:
    """Tests for GoogleAI._resolve_embedder method."""
    plugin = googleai_plugin_instance

    action = plugin._resolve_embedder(name=expected_model_name)

    assert action is not None
    assert action.kind == ActionKind.EMBEDDER
    assert action.name == expected_model_name


@pytest.mark.asyncio
async def test_googleai_list_actions(googleai_plugin_instance: GoogleAI) -> None:
    """Unit test for list actions."""

    @dataclass
    class MockModel:
        supported_actions: list[str]
        name: str
        description: str = ''

    models_return_value = [
        MockModel(supported_actions=['generateContent'], name='models/gemini-pro'),
        MockModel(supported_actions=['embedContent'], name='models/gemini-embedding-001'),
        MockModel(supported_actions=['generateContent'], name='models/gemini-2.0-flash-tts'),  # TTS
    ]

    mock_client = MagicMock()
    mock_client.models.list.return_value = models_return_value
    googleai_plugin_instance._runtime_client = lambda: mock_client

    result = await googleai_plugin_instance.list_actions()

    # Check Gemini Pro
    action1 = next(a for a in result if a.name == googleai_name('gemini-pro'))
    assert action1 is not None

    # Check Embedder
    action2 = next(a for a in result if a.name == googleai_name('gemini-embedding-001'))
    assert action2 is not None
    assert action2.action_type == ActionKind.EMBEDDER

    # Check TTS
    action3 = next(a for a in result if a.name == googleai_name('gemini-2.0-flash-tts'))
    assert action3 is not None
    # from genkit.plugins.google_genai.models.gemini import GeminiTtsConfigSchema, GeminiConfigSchema
    # assert action3.config_schema == GeminiTtsConfigSchema
    # assert action1.config_schema == GeminiConfigSchema


@pytest.mark.asyncio
async def test_googleai_list_known_models(googleai_plugin_instance: GoogleAI) -> None:
    """Unit test for list known models."""

    @dataclass
    class MockModel:
        supported_actions: list[str]
        name: str
        description: str = ''

    models_return_value = [
        MockModel(supported_actions=['generateContent'], name='models/gemini-pro'),
        MockModel(supported_actions=['embedContent'], name='models/gemini-embedding-001'),
        MockModel(supported_actions=['generateContent'], name='models/gemini-2.0-flash-tts'),  # TTS
    ]

    mock_client = MagicMock()
    mock_client.models.list.return_value = models_return_value
    googleai_plugin_instance._runtime_client = lambda: mock_client

    result = googleai_plugin_instance._list_known_models()

    # Check Gemini Pro
    action1 = next(a for a in result if a.name == googleai_name('gemini-pro'))
    assert action1 is not None

    # Check TTS
    action3 = next(a for a in result if a.name == googleai_name('gemini-2.0-flash-tts'))
    assert action3 is not None


@pytest.mark.asyncio
async def test_googleai_list_known_veo_models(googleai_plugin_instance: GoogleAI) -> None:
    """Unit test for list known veo models."""

    @dataclass
    class MockModel:
        supported_actions: list[str]
        name: str
        description: str = ''

    models_return_value = [
        MockModel(supported_actions=['generateVideos'], name='models/veo-2.0-generate-001'),
    ]

    mock_client = MagicMock()
    mock_client.models.list.return_value = models_return_value
    googleai_plugin_instance._runtime_client = lambda: mock_client

    result = googleai_plugin_instance._list_known_veo_models()

    # Check Veo
    action1 = next(a for a in result if a.name == googleai_name('veo-2.0-generate-001'))
    assert action1 is not None


@pytest.mark.asyncio
async def test_googleai_list_known_embedders(googleai_plugin_instance: GoogleAI) -> None:
    """Unit test for list known embedders."""

    @dataclass
    class MockModel:
        supported_actions: list[str]
        name: str
        description: str = ''

    models_return_value = [
        MockModel(supported_actions=['embedContent'], name='models/gemini-embedding-001'),
    ]

    mock_client = MagicMock()
    mock_client.models.list.return_value = models_return_value
    googleai_plugin_instance._runtime_client = lambda: mock_client

    result = googleai_plugin_instance._list_known_embedders()

    # Check Embedder
    action1 = next(a for a in result if a.name == googleai_name('gemini-embedding-001'))
    assert action1 is not None


@pytest.mark.parametrize(
    'input_options, expected_headers',
    [
        (
            None,
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {},
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {'headers': {'existing-header': 'value'}},
            {
                'existing-header': 'value',
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {'headers': {'x-goog-api-client': 'initial-client'}},
            {
                'x-goog-api-client': f'initial-client {GENKIT_CLIENT_HEADER}',
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {'headers': {'user-agent': 'initial-agent'}},
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': f'initial-agent {GENKIT_CLIENT_HEADER}',
            },
        ),
        (
            {'headers': {'x-goog-api-client': 'old', 'user-agent': 'old-ua'}},
            {
                'x-goog-api-client': f'old {GENKIT_CLIENT_HEADER}',
                'user-agent': f'old-ua {GENKIT_CLIENT_HEADER}',
            },
        ),
        (
            HttpOptions(),
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            HttpOptions(headers={'other': 'value'}),
            {
                'other': 'value',
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            HttpOptions(headers={'x-goog-api-client': 'initial'}),
            {
                'x-goog-api-client': f'initial {GENKIT_CLIENT_HEADER}',
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            HttpOptions(headers={'user-agent': 'initial-u'}),
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': f'initial-u {GENKIT_CLIENT_HEADER}',
            },
        ),
        (
            HttpOptions(headers={'x-goog-api-client': 'pre', 'user-agent': 'pre-u'}),
            {
                'x-goog-api-client': f'pre {GENKIT_CLIENT_HEADER}',
                'user-agent': f'pre-u {GENKIT_CLIENT_HEADER}',
            },
        ),
        (
            HttpOptions(timeout=10),
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {'timeout': 5},
            {
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
        (
            {'headers': {'one': '1'}},
            {
                'one': '1',
                'x-goog-api-client': GENKIT_CLIENT_HEADER,
                'user-agent': GENKIT_CLIENT_HEADER,
            },
        ),
    ],
)
def test_inject_attribution_headers(
    input_options: HttpOptions | dict[str, object] | None, expected_headers: dict[str, str]
) -> None:
    """Tests the _inject_attribution_headers function with various inputs."""
    result = _inject_attribution_headers(input_options)  # type: ignore
    assert isinstance(result, HttpOptions)
    assert result.headers == expected_headers


class TestVertexAIInit(unittest.TestCase):
    """Test cases for VertexAI.__init__ plugin."""

    @patch('google.genai.client.Client')
    @patch.dict(os.environ, {'GCLOUD_PROJECT': 'project'})
    def test_init_with_api_key(self, mock_genai_client: MagicMock) -> None:
        """Test using api_key parameter."""
        api_key = 'test_api_key'
        plugin = VertexAI(api_key=api_key)
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=True,
            api_key=api_key,
            credentials=None,
            project='project',
            location='us-central1',
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, VertexAI)
        self.assertTrue(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)

    @patch('google.genai.client.Client')
    @patch.dict(os.environ, {'GCLOUD_PROJECT': 'project'})
    def test_init_with_credentials(self, mock_genai_client: MagicMock) -> None:
        """Test using credentials parameter."""
        mock_credentials = MagicMock(spec=Credentials)
        plugin = VertexAI(credentials=mock_credentials)
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=True,
            api_key=None,
            credentials=mock_credentials,
            project='project',
            location='us-central1',
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, VertexAI)
        self.assertTrue(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)

    @patch('google.genai.client.Client')
    def test_init_with_all(self, mock_genai_client: MagicMock) -> None:
        """Test using credentials parameter."""
        mock_credentials = MagicMock(spec=Credentials)
        api_key = 'test_api_key'
        plugin = VertexAI(
            credentials=mock_credentials,
            api_key=api_key,
            project='project',
            location='location',
        )
        runtime_client = asyncio.run(_get_runtime_client(plugin))
        mock_genai_client.assert_called_once_with(
            vertexai=True,
            api_key=api_key,
            credentials=mock_credentials,
            project='project',
            location='location',
            debug_config=None,
            http_options=_inject_attribution_headers(),
        )
        self.assertIsInstance(plugin, VertexAI)
        self.assertTrue(plugin._vertexai)
        self.assertIsInstance(runtime_client, MagicMock)


@pytest.fixture
@patch('google.genai.client.Client')
def vertexai_plugin_instance(client: MagicMock) -> VertexAI:
    """VertexAI fixture."""
    return VertexAI(project='test-project', location='us-central1')


@pytest.mark.asyncio
async def test_vertexai_initialize(vertexai_plugin_instance: VertexAI) -> None:
    """Unit tests for VertexAI.init method."""
    plugin = vertexai_plugin_instance

    # Configure mock client to return models
    m1 = MagicMock()
    m1.name = 'publishers/google/models/gemini-1.5-flash'
    m1.supported_actions = ['generateContent']

    m2 = MagicMock()
    m2.name = 'publishers/google/models/gemini-embedding-001'
    m2.supported_actions = ['embedContent']

    mock_client = MagicMock()
    mock_client.models.list.return_value = [m1, m2]
    plugin._runtime_client = lambda: mock_client

    await plugin.init()

    # init returns known models and embedders in internal registry, but list_actions returns them list
    result = await plugin.list_actions()

    assert len(result) > 0, 'Should initialize with known models and embedders'
    assert all(action.action_type is not None for action in result), 'All actions should have an action_type'

    # ... (rest of test unchanged)

    assert all(hasattr(action, 'name') for action in result), 'All actions should have a name'
    assert all(action.name.startswith('vertexai/') for action in result), (
        "All actions should be namespaced with 'vertexai/'"
    )

    # Verify we have both models and embedders
    model_actions = [a for a in result if a.action_type == ActionKind.MODEL]
    embedder_actions = [a for a in result if a.action_type == ActionKind.EMBEDDER]
    assert len(model_actions) > 0, 'Should have at least one model'
    assert len(embedder_actions) > 0, 'Should have at least one embedder'


@patch('genkit.plugins.google_genai.VertexAI._resolve_model')
@pytest.mark.asyncio
async def test_vertexai_resolve_action_model(
    mock_resolve_action: MagicMock, vertexai_plugin_instance: VertexAI
) -> None:
    """Test resolve action for model."""
    plugin = vertexai_plugin_instance

    await plugin.resolve(action_type=ActionKind.MODEL, name='lazaro-model')
    mock_resolve_action.assert_called_once_with('lazaro-model')


@patch('genkit.plugins.google_genai.VertexAI._resolve_embedder')
@pytest.mark.asyncio
async def test_vertexai_resolve_action_embedder(
    mock_resolve_action: MagicMock, vertexai_plugin_instance: VertexAI
) -> None:
    """Test resolve action for embedder."""
    plugin = vertexai_plugin_instance

    await plugin.resolve(action_type=ActionKind.EMBEDDER, name='lazaro-model')
    mock_resolve_action.assert_called_once_with('lazaro-model')


@patch(
    'genkit.plugins.google_genai.models.gemini.google_model_info',
    new_callable=MagicMock,
)
@patch(
    'genkit.plugins.google_genai.models.imagen.vertexai_image_model_info',
    new_callable=MagicMock,
)
@pytest.mark.parametrize(
    'model_name, expected_model_name, key, image',
    [
        (
            'gemini-pro-deluxe-max',
            'vertexai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
            False,
        ),
        (
            'vertexai/gemini-pro-deluxe-max',
            'vertexai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
            False,
        ),
        (
            'vertexai/image-gemini-pro-deluxe-max',
            'vertexai/image-gemini-pro-deluxe-max',
            'image-gemini-pro-deluxe-max',
            True,
        ),
        (
            'image-gemini-pro-deluxe-max',
            'vertexai/image-gemini-pro-deluxe-max',
            'image-gemini-pro-deluxe-max',
            True,
        ),
        (
            'gemini-pro-deluxe-max-image',
            'vertexai/gemini-pro-deluxe-max-image',
            'gemini-pro-deluxe-max-image',
            False,
        ),
    ],
)
def test_vertexai__resolve_model(
    mock_google_model_info: MagicMock,
    mock_vertexai_image_model_info: MagicMock,
    model_name: str,
    expected_model_name: str,
    key: str,
    image: bool,
    vertexai_plugin_instance: VertexAI,
) -> None:
    """Tests for VertexAI._resolve_model method."""
    plugin = vertexai_plugin_instance
    MagicMock(spec=Genkit)

    mock_google_model_info.return_value = ModelInfo(
        label=f'Google AI - {model_name}',
        supports=DEFAULT_SUPPORTS_MODEL,
    )

    mock_vertexai_image_model_info.return_value = ModelInfo(
        label=f'Vertex AI - {model_name}',
        supports=DEFAULT_IMAGE_SUPPORT,
    )

    action = plugin._resolve_model(name=expected_model_name)

    assert action is not None
    assert action.kind == ActionKind.MODEL
    assert action.name == expected_model_name

    if image:
        assert key in IMAGE_SUPPORTED_MODELS
    else:
        assert key in SUPPORTED_MODELS


@pytest.mark.parametrize(
    'model_name, expected_model_name, clean_name',
    [
        (
            'gemini-pro-deluxe-max',
            'vertexai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
        ),
        (
            'vertexai/gemini-pro-deluxe-max',
            'vertexai/gemini-pro-deluxe-max',
            'gemini-pro-deluxe-max',
        ),
    ],
)
def test_vertexai__resolve_embedder(
    model_name: str,
    expected_model_name: str,
    clean_name: str,
    vertexai_plugin_instance: VertexAI,
) -> None:
    """Tests for VertexAI._resolve_embedder method."""
    plugin = vertexai_plugin_instance

    action = plugin._resolve_embedder(name=expected_model_name)

    assert action is not None
    assert action.kind == ActionKind.EMBEDDER
    assert action.name == expected_model_name


@pytest.mark.asyncio
async def test_vertexai_list_actions(vertexai_plugin_instance: VertexAI) -> None:
    """Unit test for list actions."""

    @dataclass
    class MockModel:
        name: str
        description: str = ''

    [
        MockModel(name='publishers/google/models/gemini-1.5-flash'),
        MockModel(name='publishers/google/models/gemini-embedding-001'),
        MockModel(name='publishers/google/models/imagen-3.0-generate-001'),
        MockModel(name='publishers/google/models/veo-2.0-generate-001'),
    ]

    mock_client = MagicMock()
    # Create sophisticated mocks that have supported_actions
    m1 = MagicMock()
    m1.name = 'publishers/google/models/gemini-1.5-flash'
    m1.supported_actions = ['generateContent']
    m1.description = 'Gemini model'

    m2 = MagicMock()
    m2.name = 'publishers/google/models/gemini-embedding-001'
    m2.supported_actions = ['embedContent']
    m2.description = 'Embedder'

    m3 = MagicMock()
    m3.name = 'publishers/google/models/imagen-3.0-generate-001'
    m3.supported_actions = ['predict']  # Imagen uses predict
    m3.description = 'Imagen'

    m4 = MagicMock()
    m4.name = 'publishers/google/models/veo-2.0-generate-001'
    m4.supported_actions = ['generateVideos']  # Veo uses generateVideos
    m4.description = 'Veo'

    mock_client.models.list.return_value = [m1, m2, m3, m4]
    vertexai_plugin_instance._runtime_client = lambda: mock_client

    result = await vertexai_plugin_instance.list_actions()

    # Verify Gemini
    action1 = next(a for a in result if a.name == vertexai_name('gemini-1.5-flash'))
    assert action1 is not None

    # Verify Embedder
    action2 = next(a for a in result if a.name == vertexai_name('gemini-embedding-001'))
    assert action2 is not None

    # Verify Imagen
    action3 = next(a for a in result if a.name == vertexai_name('imagen-3.0-generate-001'))
    assert action3 is not None
    assert action3.action_type == ActionKind.MODEL

    # Verify Veo
    action4 = next(a for a in result if a.name == vertexai_name('veo-2.0-generate-001'))
    assert action4 is not None
    # from genkit.plugins.google_genai.models.veo import VeoConfigSchema
    # assert action4.config_schema == VeoConfigSchema


@pytest.mark.asyncio
async def test_vertexai_list_actions_without_supported_actions(vertexai_plugin_instance: VertexAI) -> None:
    """Regression test for #5572.

    Vertex AI's ``client.models.list()`` returns publisher models with
    ``supported_actions = None``. Discovery must categorize these by name
    rather than skipping them, otherwise no Vertex models appear in the Dev UI.
    """

    def mock_model(name: str) -> MagicMock:
        m = MagicMock()
        m.name = name
        m.supported_actions = None  # Vertex leaves this unset.
        m.description = ''
        return m

    mock_client = MagicMock()
    mock_client.models.list.return_value = [
        mock_model('publishers/google/models/gemini-2.5-pro'),
        mock_model('publishers/google/models/gemini-embedding-001'),
        mock_model('publishers/google/models/imagen-3.0-generate-002'),
        mock_model('publishers/google/models/veo-2.0-generate-001'),
    ]
    vertexai_plugin_instance._runtime_client = lambda: mock_client

    result = await vertexai_plugin_instance.list_actions()
    names = {a.name for a in result}

    # Gemini text model discovered despite supported_actions=None.
    assert vertexai_name('gemini-2.5-pro') in names
    # Imagen and Veo discovered.
    assert vertexai_name('imagen-3.0-generate-002') in names
    assert vertexai_name('veo-2.0-generate-001') in names

    # gemini-embedding-001 is registered as an embedder, not a gemini model.
    embedder = next(a for a in result if a.name == vertexai_name('gemini-embedding-001'))
    assert embedder.action_type == ActionKind.EMBEDDER


@pytest.mark.asyncio
async def test_googleai_resolve_background_model(googleai_plugin_instance: GoogleAI) -> None:
    """Test resolve action for background model."""
    plugin = googleai_plugin_instance

    action = await plugin.resolve(action_type=ActionKind.BACKGROUND_MODEL, name=googleai_name('veo-2.0-generate-001'))
    assert action is not None
    assert action.kind == ActionKind.BACKGROUND_MODEL
    assert action.name == googleai_name('veo-2.0-generate-001')


@pytest.mark.asyncio
async def test_googleai_resolve_check_operation(googleai_plugin_instance: GoogleAI) -> None:
    """Test resolve action for check operation."""
    plugin = googleai_plugin_instance

    action = await plugin.resolve(
        action_type=ActionKind.CHECK_OPERATION, name=googleai_name('veo-2.0-generate-001/check')
    )
    assert action is not None
    assert action.kind == ActionKind.CHECK_OPERATION
    assert action.name == googleai_name('veo-2.0-generate-001/check')


@pytest.mark.asyncio
async def test_vertexai_list_known_models(vertexai_plugin_instance: VertexAI) -> None:
    """Unit test for list known models."""

    @dataclass
    class MockModel:
        name: str
        description: str = ''

    [
        MockModel(name='publishers/google/models/gemini-1.5-flash'),
        MockModel(name='publishers/google/models/gemini-embedding-001'),
        MockModel(name='publishers/google/models/imagen-3.0-generate-001'),
        MockModel(name='publishers/google/models/veo-2.0-generate-001'),
    ]

    mock_client = MagicMock()
    # Create sophisticated mocks that have supported_actions
    m1 = MagicMock()
    m1.name = 'publishers/google/models/gemini-1.5-flash'
    m1.supported_actions = ['generateContent']
    m1.description = 'Gemini model'

    m2 = MagicMock()
    m2.name = 'publishers/google/models/gemini-embedding-001'
    m2.supported_actions = ['embedContent']
    m2.description = 'Embedder'

    m3 = MagicMock()
    m3.name = 'publishers/google/models/imagen-3.0-generate-001'
    m3.supported_actions = ['predict']
    m3.description = 'Imagen'

    m4 = MagicMock()
    m4.name = 'publishers/google/models/veo-2.0-generate-001'
    m4.supported_actions = ['generateVideos']
    m4.description = 'Veo'

    mock_client.models.list.return_value = [m1, m2, m3, m4]
    vertexai_plugin_instance._runtime_client = lambda: mock_client

    result = vertexai_plugin_instance._list_known_models()

    # Verify Gemini
    action1 = next(a for a in result if a.name == vertexai_name('gemini-1.5-flash'))
    assert action1 is not None

    # Verify Imagen
    action3 = next(a for a in result if a.name == vertexai_name('imagen-3.0-generate-001'))
    assert action3 is not None

    # Verify Veo
    action4 = next(a for a in result if a.name == vertexai_name('veo-2.0-generate-001'))
    assert action4 is not None


@pytest.mark.asyncio
async def test_vertexai_list_known_embedders(vertexai_plugin_instance: VertexAI) -> None:
    """Unit test for list known embedders."""

    @dataclass
    class MockModel:
        name: str
        description: str = ''

    [
        MockModel(name='publishers/google/models/gemini-1.5-flash'),
        MockModel(name='publishers/google/models/gemini-embedding-001'),
        MockModel(name='publishers/google/models/imagen-3.0-generate-001'),
        MockModel(name='publishers/google/models/veo-2.0-generate-001'),
    ]

    mock_client = MagicMock()
    # Create sophisticated mocks that have supported_actions
    m1 = MagicMock()
    m1.name = 'publishers/google/models/gemini-1.5-flash'
    m1.supported_actions = ['generateContent']
    m1.description = 'Gemini model'

    m2 = MagicMock()
    m2.name = 'publishers/google/models/gemini-embedding-001'
    m2.supported_actions = ['embedContent']
    m2.description = 'Embedder'

    m3 = MagicMock()
    m3.name = 'publishers/google/models/imagen-3.0-generate-001'
    m3.supported_actions = ['predict']
    m3.description = 'Imagen'

    m4 = MagicMock()
    m4.name = 'publishers/google/models/veo-2.0-generate-001'
    m4.supported_actions = ['generateVideos']
    m4.description = 'Veo'

    mock_client.models.list.return_value = [m1, m2, m3, m4]
    vertexai_plugin_instance._runtime_client = lambda: mock_client

    result = vertexai_plugin_instance._list_known_embedders()

    # Verify Embedder
    action2 = next(a for a in result if a.name == vertexai_name('gemini-embedding-001'))
    assert action2 is not None


@pytest.mark.asyncio
async def test_vertexai_resolve_evaluator(vertexai_plugin_instance: VertexAI) -> None:
    """Test resolve action for evaluator."""
    plugin = vertexai_plugin_instance

    action = await plugin.resolve(action_type=ActionKind.EVALUATOR, name=vertexai_name('fluency'))
    assert action is not None
    assert action.kind == ActionKind.EVALUATOR
    assert action.name == vertexai_name('fluency')


def test_config_schema_extra_fields() -> None:
    """Test that config schema accepts extra fields (dynamic config)."""
    # Validation should succeed with unknown field by using model_validate for dynamic fields
    # to avoid static type checker errors on constructor
    config_data = {'temperature': 0.5, 'new_experimental_param': 'test'}
    config = GeminiConfigSchema.model_validate(config_data)

    assert config.temperature == 0.5
    # Access dynamic fields via getattr or __dict__ to make type checker happy
    assert config.new_experimental_param == 'test'  # type: ignore
    assert config.model_dump()['new_experimental_param'] == 'test'


@pytest.mark.asyncio
async def test_system_prompt_handling() -> None:
    """Test that system prompts are correctly extracted to config."""
    mock_client = MagicMock(spec=genai.Client)
    model = GeminiModel(version='gemini-1.5-flash', client=mock_client)

    request = ModelRequest(
        messages=[
            Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='You are a helpful assistant'))]),
            Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]),
        ],
        config=None,
    )

    cfg = await model._genkit_to_googleai_cfg(request)

    assert cfg is not None
    assert cfg.system_instruction is not None
    assert cfg.system_instruction.parts is not None  # type: ignore
    assert len(cfg.system_instruction.parts) == 1  # type: ignore
    assert cfg.system_instruction.parts[0].text == 'You are a helpful assistant'  # type: ignore
