# 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.

"""Tests for model name utility functions."""

from google.adk.models.llm_request import LlmRequest
from google.adk.utils.model_name_utils import _is_gemini_3_x_live
from google.adk.utils.model_name_utils import _is_managed_agent
from google.adk.utils.model_name_utils import extract_model_name
from google.adk.utils.model_name_utils import is_gemini_1_model
from google.adk.utils.model_name_utils import is_gemini_3_5_live_translate
from google.adk.utils.model_name_utils import is_gemini_eap_or_2_or_above
from google.adk.utils.model_name_utils import is_gemini_model
from google.adk.utils.model_name_utils import is_gemini_model_id_check_disabled


class TestExtractModelName:
  """Test the extract_model_name function."""

  def test_extract_model_name_simple_model(self):
    """Test extraction of simple model names."""
    assert extract_model_name('gemini-2.5-pro') == 'gemini-2.5-pro'
    assert extract_model_name('gemini-2.5-flash') == 'gemini-2.5-flash'
    assert extract_model_name('gemini-1.0-pro') == 'gemini-1.0-pro'
    assert extract_model_name('claude-3-sonnet') == 'claude-3-sonnet'
    assert extract_model_name('gpt-4') == 'gpt-4'

  def test_extract_model_name_path_based_model(self):
    """Test extraction of path-based model names."""
    path_model = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
    assert extract_model_name(path_model) == 'gemini-2.5-flash'

    path_model_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
    assert extract_model_name(path_model_2) == 'gemini-1.5-pro-preview'

    path_model_3 = 'projects/test-project/locations/europe-west1/publishers/google/models/claude-3-sonnet'
    assert extract_model_name(path_model_3) == 'claude-3-sonnet'

    path_model_4 = 'apigee/gemini-2.5-flash'
    assert extract_model_name(path_model_4) == 'gemini-2.5-flash'

    path_model_5 = 'apigee/v1/gemini-2.5-flash'
    assert extract_model_name(path_model_5) == 'gemini-2.5-flash'

    path_model_6 = 'apigee/gemini/gemini-2.5-flash'
    assert extract_model_name(path_model_6) == 'gemini-2.5-flash'

    path_model_7 = 'apigee/vertex_ai/gemini-2.5-flash'
    assert extract_model_name(path_model_7) == 'gemini-2.5-flash'

    path_model_8 = 'apigee/gemini/v1/gemini-2.5-flash'
    assert extract_model_name(path_model_8) == 'gemini-2.5-flash'

    path_model_9 = 'apigee/vertex_ai/v1beta/gemini-2.5-flash'
    assert extract_model_name(path_model_9) == 'gemini-2.5-flash'

  def test_extract_model_name_with_models_prefix(self):
    """Test extraction of model names with 'models/' prefix."""
    assert extract_model_name('models/gemini-2.5-pro') == 'gemini-2.5-pro'
    assert extract_model_name('models/gemini-2.5-flash') == 'gemini-2.5-flash'

  def test_extract_model_name_provider_prefixed_model(self):
    """Test extraction of provider-prefixed Gemini model names."""
    assert extract_model_name('gemini/gemini-2.5-flash') == 'gemini-2.5-flash'
    assert extract_model_name('vertex_ai/gemini-2.5-flash') == (
        'gemini-2.5-flash'
    )
    assert (
        extract_model_name('openrouter/google/gemini-2.5-pro:online')
        == 'gemini-2.5-pro:online'
    )
    assert extract_model_name('openrouter/anthropic/claude-sonnet-4') == (
        'openrouter/anthropic/claude-sonnet-4'
    )

  def test_extract_model_name_invalid_path(self):
    """Test that invalid path formats return the original string."""
    invalid_paths = [
        'projects/invalid/path/format',
        'invalid/path/format',
        'projects/123/locations/us-central1/models/gemini-2.5-flash',  # missing publishers
        'projects/123/publishers/google/models/gemini-2.5-flash',  # missing locations
        'projects/123/locations/us-central1/publishers/google/gemini-2.5-flash',  # missing models
    ]

    for invalid_path in invalid_paths:
      assert extract_model_name(invalid_path) == invalid_path

  def test_extract_model_name_empty_string(self):
    """Test extraction from empty string."""
    assert extract_model_name('') == ''

  def test_extract_model_name_edge_cases(self):
    """Test edge cases for model name extraction."""
    # Test with unusual but valid path patterns
    path_with_numbers = 'projects/123456789/locations/us-central1/publishers/google/models/gemini-2.5-flash'
    assert extract_model_name(path_with_numbers) == 'gemini-2.5-flash'

    # Test with hyphens in project/location names
    path_with_hyphens = 'projects/my-test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro'
    assert extract_model_name(path_with_hyphens) == 'gemini-1.5-pro'


class TestIsGeminiModel:
  """Test the is_gemini_model function."""

  def test_is_gemini_model_simple_names(self):
    """Test Gemini model detection with simple model names."""
    assert is_gemini_model('gemini-2.5-pro') is True
    assert is_gemini_model('gemini-1.5-flash') is True
    assert is_gemini_model('gemini-1.0-pro') is True
    assert is_gemini_model('gemini-2.5-flash') is True
    assert is_gemini_model('claude-3-sonnet') is False
    assert is_gemini_model('gpt-4') is False
    assert is_gemini_model('llama-2') is False

  def test_is_gemini_model_path_based_names(self):
    """Test Gemini model detection with path-based model names."""
    gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
    assert is_gemini_model(gemini_path) is True

    gemini_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
    assert is_gemini_model(gemini_path_2) is True

    non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
    assert is_gemini_model(non_gemini_path) is False

  def test_is_gemini_model_provider_prefixed_names(self):
    """Test Gemini model detection with provider-prefixed model names."""
    assert is_gemini_model('gemini/gemini-2.5-flash') is True
    assert is_gemini_model('vertex_ai/gemini-2.5-flash') is True
    assert is_gemini_model('openrouter/google/gemini-2.5-pro:online') is True
    assert is_gemini_model('openrouter/anthropic/claude-sonnet-4') is False

  def test_is_gemini_model_edge_cases(self):
    """Test edge cases for Gemini model detection."""
    # Test with None
    assert is_gemini_model(None) is False

    # Test with empty string
    assert is_gemini_model('') is False

    # Test with model names containing gemini but not starting with it
    assert is_gemini_model('my-gemini-model') is False
    assert is_gemini_model('claude-gemini-hybrid') is False

    # Test with model names that have gemini in the middle of the path
    tricky_path = 'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet'
    assert is_gemini_model(tricky_path) is False

    # Test with just "gemini" without dash
    assert is_gemini_model('gemini') is False
    assert is_gemini_model('gemini_1_5_flash') is False

  def test_is_gemini_model_case_sensitivity(self):
    """Test that model detection is case-sensitive."""
    assert is_gemini_model('Gemini-2.5-pro') is False
    assert is_gemini_model('GEMINI-2.5-pro') is False
    assert is_gemini_model('gemini-2.5-PRO') is True  # Only the start matters


class TestIsGemini1Model:
  """Test the is_gemini_1_model function."""

  def test_is_gemini_1_model_simple_names(self):
    """Test Gemini 1.x model detection with simple model names."""
    assert is_gemini_1_model('gemini-1.5-flash') is True
    assert is_gemini_1_model('gemini-1.0-pro') is True
    assert is_gemini_1_model('gemini-1.5-pro-preview') is True
    assert is_gemini_1_model('gemini-1.9-experimental') is True
    assert is_gemini_1_model('gemini-2.5-flash') is False
    assert is_gemini_1_model('gemini-2.5-pro') is False
    assert is_gemini_1_model('gemini-10.0-pro') is False  # Only 1.x versions
    assert is_gemini_1_model('claude-3-sonnet') is False

  def test_is_gemini_1_model_path_based_names(self):
    """Test Gemini 1.x model detection with path-based model names."""
    gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash'
    assert is_gemini_1_model(gemini_1_path) is True

    gemini_1_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.0-pro-preview'
    assert is_gemini_1_model(gemini_1_path_2) is True

    gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
    assert is_gemini_1_model(gemini_2_path) is False

  def test_is_gemini_1_model_provider_prefixed_names(self):
    """Test Gemini 1.x detection with provider-prefixed model names."""
    assert is_gemini_1_model('gemini/gemini-1.5-flash') is True
    assert is_gemini_1_model('vertex_ai/gemini-1.5-flash') is True
    assert is_gemini_1_model('openrouter/google/gemini-1.5-pro:online') is True
    assert is_gemini_1_model('openrouter/google/gemini-2.5-pro') is False

  def test_is_gemini_1_model_edge_cases(self):
    """Test edge cases for Gemini 1.x model detection."""
    # Test with None
    assert is_gemini_1_model(None) is False

    # Test with empty string
    assert is_gemini_1_model('') is False

    # Test with model names containing gemini-1 but not starting with it
    assert is_gemini_1_model('my-gemini-1.5-model') is False
    assert is_gemini_1_model('custom-gemini-1.5-flash') is False

    # Test with invalid versions
    assert is_gemini_1_model('gemini-1') is False  # Missing dot
    assert is_gemini_1_model('gemini-1-pro') is False  # Missing dot
    assert is_gemini_1_model('gemini-1.') is False  # Missing version number


class TestIsGemini2Model:
  """Test the is_gemini_eap_or_2_or_above function."""

  def test_is_gemini_eap_or_2_or_above_simple_names(self):
    """Test Gemini 2.0+ model detection with simple model names."""
    assert is_gemini_eap_or_2_or_above('gemini-2.5-flash') is True
    assert is_gemini_eap_or_2_or_above('gemini-2.5-pro') is True
    assert is_gemini_eap_or_2_or_above('gemini-2.9-experimental') is True
    assert is_gemini_eap_or_2_or_above('gemini-2-pro') is True
    assert is_gemini_eap_or_2_or_above('gemini-2') is True
    assert is_gemini_eap_or_2_or_above('gemini-3.0-pro') is True
    assert is_gemini_eap_or_2_or_above('gemini-flash-early-exp') is True
    assert is_gemini_eap_or_2_or_above('gemini-flash-early-exp3') is True
    assert is_gemini_eap_or_2_or_above('gemini-flash-lite-early-exp') is True
    assert is_gemini_eap_or_2_or_above('gemini-pro-early-exp') is True
    assert is_gemini_eap_or_2_or_above('gemini-1.5-flash') is False
    assert is_gemini_eap_or_2_or_above('gemini-1.0-pro') is False
    assert is_gemini_eap_or_2_or_above('claude-3-sonnet') is False

  def test_is_gemini_eap_or_2_or_above_path_based_names(self):
    """Test Gemini 2.0+ model detection with path-based model names."""
    gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
    assert is_gemini_eap_or_2_or_above(gemini_2_path) is True

    gemini_2_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-2.5-pro-preview'
    assert is_gemini_eap_or_2_or_above(gemini_2_path_2) is True

    gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash'
    assert is_gemini_eap_or_2_or_above(gemini_1_path) is False

    gemini_3_path = 'projects/12345/locations/us-east1/publishers/google/models/gemini-3.0-pro'
    assert is_gemini_eap_or_2_or_above(gemini_3_path) is True

  def test_is_gemini_eap_or_2_or_above_provider_prefixed_names(self):
    """Test Gemini 2.0+ detection with provider-prefixed model names."""
    assert is_gemini_eap_or_2_or_above('gemini/gemini-2.5-flash') is True
    assert is_gemini_eap_or_2_or_above('vertex_ai/gemini-2.5-flash') is True
    assert (
        is_gemini_eap_or_2_or_above('openrouter/google/gemini-2.5-pro:online')
        is True
    )
    assert (
        is_gemini_eap_or_2_or_above('openrouter/google/gemini-1.5-pro:online')
        is False
    )

  def test_is_gemini_eap_or_2_or_above_edge_cases(self):
    """Test edge cases for Gemini 2.0+ model detection."""
    # Test with None
    assert is_gemini_eap_or_2_or_above(None) is False

    # Test with empty string
    assert is_gemini_eap_or_2_or_above('') is False

    # Test with model names containing gemini-2 but not starting with it
    assert is_gemini_eap_or_2_or_above('my-gemini-2.5-model') is False
    assert is_gemini_eap_or_2_or_above('custom-gemini-2.5-flash') is False

    # Test with invalid versions
    assert (
        is_gemini_eap_or_2_or_above('gemini-2.') is False
    )  # Missing version number
    assert is_gemini_eap_or_2_or_above('gemini-0.9-test') is False
    assert is_gemini_eap_or_2_or_above('gemini-one') is False


class TestModelNameUtilsIntegration:
  """Integration tests for model name utilities."""

  def test_model_classification_consistency(self):
    """Test that model classification functions are consistent."""
    test_models = [
        'gemini-1.5-flash',
        'gemini-2.5-flash',
        'gemini-2.5-pro',
        'gemini-3.0-pro',
        'gemini/gemini-2.5-flash',
        'openrouter/google/gemini-2.5-pro:online',
        'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-pro',
        'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-flash',
        'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
        'claude-3-sonnet',
        'gpt-4',
    ]

    for model in test_models:
      # A model can only be either Gemini 1.x or Gemini 2.0+, not both
      if is_gemini_1_model(model):
        assert not is_gemini_eap_or_2_or_above(
            model
        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
        assert is_gemini_model(
            model
        ), f'Model {model} is Gemini 1.x but not classified as Gemini'

      if is_gemini_eap_or_2_or_above(model):
        assert not is_gemini_1_model(
            model
        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
        assert is_gemini_model(
            model
        ), f'Model {model} is Gemini 2.0+ but not classified as Gemini'

      # If it's neither Gemini 1.x nor 2.0+, it should not be classified as Gemini
      if not is_gemini_1_model(model) and not is_gemini_eap_or_2_or_above(
          model
      ):
        if model and 'gemini-' not in extract_model_name(model):
          assert not is_gemini_model(
              model
          ), f'Non-Gemini model {model} classified as Gemini'

  def test_path_vs_simple_model_consistency(self):
    """Test that path-based and simple model names are classified consistently."""
    model_pairs = [
        (
            'gemini-1.5-flash',
            'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-flash',
        ),
        (
            'gemini-2.5-flash',
            'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-flash',
        ),
        (
            'gemini-2.5-pro',
            'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-pro',
        ),
        (
            'gemini-3.0-pro',
            'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
        ),
        (
            'claude-3-sonnet',
            'projects/123/locations/us-central1/publishers/google/models/claude-3-sonnet',
        ),
    ]

    for simple_model, path_model in model_pairs:
      # Both forms should be classified identically
      assert is_gemini_model(simple_model) == is_gemini_model(path_model), (
          f'Inconsistent Gemini classification for {simple_model} vs'
          f' {path_model}'
      )
      assert is_gemini_1_model(simple_model) == is_gemini_1_model(path_model), (
          f'Inconsistent Gemini 1.x classification for {simple_model} vs'
          f' {path_model}'
      )
      assert is_gemini_eap_or_2_or_above(
          simple_model
      ) == is_gemini_eap_or_2_or_above(path_model), (
          f'Inconsistent Gemini 2.0+ classification for {simple_model} vs'
          f' {path_model}'
      )


class TestGeminiModelIdCheckFlag:
  """Tests for Gemini model-id check override flag."""

  def test_default_is_disabled(self, monkeypatch):
    monkeypatch.delenv('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', raising=False)
    assert is_gemini_model_id_check_disabled() is False

  def test_true_enables_check_bypass(self, monkeypatch):
    monkeypatch.setenv('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', 'true')
    assert is_gemini_model_id_check_disabled() is True


class TestIsGemini3XLive:
  """Test the _is_gemini_3_x_live function."""

  def test_is_gemini_3_x_live_simple_name(self):
    """Test with simple model name format."""
    assert _is_gemini_3_x_live('gemini-3.1-flash-live') is True
    assert _is_gemini_3_x_live('gemini-3.1-flash-live-preview') is True
    assert _is_gemini_3_x_live('gemini-3.5-flash-lite-live-preview') is True
    assert _is_gemini_3_x_live('gemini-3.5-live-translate') is False
    assert _is_gemini_3_x_live('gemini-3.1-pro') is False
    assert _is_gemini_3_x_live('gemini-2.5-flash-live') is False

  def test_is_gemini_3_x_live_path_based_name(self):
    """Test with path-based format (Vertex AI etc.)."""
    vertex_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash-live'
    assert _is_gemini_3_x_live(vertex_path) is True

    vertex_path_preview = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite-live-preview'
    assert _is_gemini_3_x_live(vertex_path_preview) is True

    non_live_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash'
    assert _is_gemini_3_x_live(non_live_path) is False

  def test_is_gemini_3_x_live_edge_cases(self):
    """Test edge cases."""
    assert _is_gemini_3_x_live(None) is False
    assert _is_gemini_3_x_live('') is False


class TestIsGemini35LiveTranslate:
  """Test the is_gemini_3_5_live_translate function."""

  def test_is_gemini_3_5_live_translate_simple_name(self):
    """Test with simple model name format."""
    assert is_gemini_3_5_live_translate('gemini-3.5-live-translate') is True
    assert is_gemini_3_5_live_translate('gemini-3.5-flash-live') is False

  def test_is_gemini_3_5_live_translate_path_based_name(self):
    """Test with path-based format (Vertex AI etc.)."""
    vertex_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.5-live-translate-preview'
    assert is_gemini_3_5_live_translate(vertex_path) is True

  def test_is_gemini_3_5_live_translate_edge_cases(self):
    """Test edge cases."""
    assert is_gemini_3_5_live_translate(None) is False
    assert is_gemini_3_5_live_translate('') is False


class TestIsManagedAgent:
  """Tests for the _is_managed_agent predicate."""

  def test_true_when_flag_set(self):
    request = LlmRequest()
    request._is_managed_agent = True
    assert _is_managed_agent(request) is True

  def test_false_by_default(self):
    assert _is_managed_agent(LlmRequest()) is False