import os
import sys
from unittest.mock import patch

import pytest

from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH

sys.path.insert(
    0, os.path.abspath("../../..")
)  # Adds the parent directory to the system path

from litellm.llms.vertex_ai.common_utils import (
    _get_vertex_url,
    convert_anyof_null_to_nullable,
    get_vertex_location_from_url,
    get_vertex_project_id_from_url,
    pop_vertex_request_labels,
    set_schema_property_ordering,
    vertex_request_labels_from_litellm_params,
)


@pytest.mark.asyncio
async def test_get_vertex_project_id_from_url():
    """Test _get_vertex_project_id_from_url with various URLs"""
    # Test with valid URL
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
    project_id = get_vertex_project_id_from_url(url)
    assert project_id == "test-project"

    # Test with invalid URL
    url = "https://invalid-url.com"
    project_id = get_vertex_project_id_from_url(url)
    assert project_id is None


@pytest.mark.asyncio
async def test_get_vertex_location_from_url():
    """Test _get_vertex_location_from_url with various URLs"""
    # Test with valid URL
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
    location = get_vertex_location_from_url(url)
    assert location == "us-central1"

    # Test with invalid URL
    url = "https://invalid-url.com"
    location = get_vertex_location_from_url(url)
    assert location is None


def test_basic_anyof_conversion():
    """Test basic conversion of anyOf with 'null'."""
    schema = {
        "type": "object",
        "properties": {"example": {"anyOf": [{"type": "string"}, {"type": "null"}]}},
    }

    convert_anyof_null_to_nullable(schema)

    expected = {
        "type": "object",
        "properties": {"example": {"anyOf": [{"type": "string", "nullable": True}]}},
    }
    assert schema == expected


def test_nested_anyof_conversion():
    """Test nested conversion with 'anyOf' inside properties."""
    schema = {
        "type": "object",
        "properties": {
            "outer": {
                "type": "object",
                "properties": {
                    "inner": {
                        "anyOf": [
                            {"type": "array", "items": {"type": "string"}},
                            {"type": "string"},
                            {"type": "null"},
                        ]
                    }
                },
            }
        },
    }

    convert_anyof_null_to_nullable(schema)

    expected = {
        "type": "object",
        "properties": {
            "outer": {
                "type": "object",
                "properties": {
                    "inner": {
                        "anyOf": [
                            {
                                "type": "array",
                                "items": {"type": "string"},
                                "nullable": True,
                            },
                            {"type": "string", "nullable": True},
                        ]
                    }
                },
            }
        },
    }
    assert schema == expected


def test_anyof_with_excessive_nesting():
    """Test conversion with excessive nesting > max levels +1 deep."""
    # generate a schema with excessive nesting
    schema = {"type": "object", "properties": {}}
    current = schema
    for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
        current["properties"] = {
            "nested": {
                "anyOf": [{"type": "string"}, {"type": "null"}],
                "properties": {},
            }
        }
        current = current["properties"]["nested"]

    # running the conversion will raise an error
    with pytest.raises(
        ValueError,
        match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting.",
    ):
        convert_anyof_null_to_nullable(schema)


@pytest.mark.asyncio
async def test_get_supports_system_message():
    """Test get_supports_system_message with different models"""
    from litellm.llms.vertex_ai.common_utils import get_supports_system_message

    # fine-tuned vertex gemini models will specifiy they are in the /gemini spec format
    result = get_supports_system_message(
        model="gemini/1234567890", custom_llm_provider="vertex_ai"
    )
    assert result == True

    # non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format
    result = get_supports_system_message(
        model="random-model-name", custom_llm_provider="vertex_ai"
    )
    assert result == False


def test_set_schema_property_ordering_with_excessive_nesting():
    """Test set_schema_property_ordering with excessive nesting > max levels +1 deep."""
    # generate a schema with excessive nesting
    schema = {"type": "object", "properties": {}}
    current = schema
    for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
        current["properties"] = {"nested": {"type": "object", "properties": {}}}
        current = current["properties"]["nested"]

    # running the function will raise an error
    with pytest.raises(
        ValueError,
        match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting.",
    ):
        set_schema_property_ordering(schema)


def test_build_vertex_schema():
    """Test build_vertex_schema with a sample schema"""
    from litellm.llms.vertex_ai.common_utils import _build_vertex_schema

    parameters = {
        "properties": {
            "state": {
                "properties": {
                    "messages": {"items": {}, "type": "array"},
                    "conversation_id": {"type": "string"},
                },
                "required": ["messages", "conversation_id"],
                "type": "object",
            },
            "config": {
                "description": "Configuration for a Runnable.",
                "properties": {
                    "tags": {"items": {"type": "string"}, "type": "array"},
                    "metadata": {"type": "object"},
                    "callbacks": {
                        "anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]
                    },
                    "run_name": {"type": "string"},
                    "max_concurrency": {
                        "anyOf": [{"type": "integer"}, {"type": "null"}]
                    },
                    "recursion_limit": {"type": "integer"},
                    "configurable": {"type": "object"},
                    "run_id": {
                        "anyOf": [
                            {"format": "uuid", "type": "string"},
                            {"type": "null"},
                        ]
                    },
                },
                "type": "object",
            },
            "kwargs": {"default": None, "type": "object"},
        },
        "required": ["state", "config"],
        "type": "object",
    }

    expected_output = {
        "properties": {
            "state": {
                "properties": {
                    "messages": {"items": {"type": "object"}, "type": "array"},
                    "conversation_id": {"type": "string"},
                },
                "required": ["messages", "conversation_id"],
                "type": "object",
            },
            "config": {
                "description": "Configuration for a Runnable.",
                "properties": {
                    "tags": {"items": {"type": "string"}, "type": "array"},
                    "metadata": {"type": "object"},
                    "callbacks": {
                        "anyOf": [
                            {
                                "type": "array",
                                "items": {"type": "object"},
                                "nullable": True,
                            },
                            {"type": "object", "nullable": True},
                        ]
                    },
                    "run_name": {"type": "string"},
                    "max_concurrency": {
                        "anyOf": [{"type": "integer", "nullable": True}]
                    },
                    "recursion_limit": {"type": "integer"},
                    "configurable": {"type": "object"},
                    "run_id": {"anyOf": [{"type": "string", "nullable": True}]},
                },
                "type": "object",
            },
            "kwargs": {"default": None, "type": "object"},
        },
        "required": ["state", "config"],
        "type": "object",
    }

    assert _build_vertex_schema(parameters) == expected_output


def test_process_items_with_excessive_nesting():
    """Test process_items with excessive nesting > max levels +1 deep."""
    # generate a schema with excessive nesting
    from litellm.llms.vertex_ai.common_utils import process_items

    schema = {"type": "object", "properties": {}}
    current = schema
    for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
        current["properties"] = {"nested": {"type": "object", "properties": {}}}
        current = current["properties"]["nested"]

    with pytest.raises(
        ValueError,
        match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting.",
    ):
        process_items(schema)


def test_process_items_basic():
    """Test basic functionality of process_items."""
    from litellm.llms.vertex_ai.common_utils import process_items

    # Test empty items
    schema = {"type": "array", "items": {}}
    process_items(schema)
    assert schema["items"] == {"type": "object"}

    # Test nested items
    schema = {"type": "array", "items": {"type": "array", "items": {}}}
    process_items(schema)
    assert schema["items"]["items"] == {"type": "object"}

    # Test items in properties
    schema = {
        "type": "object",
        "properties": {"nested": {"type": "array", "items": {}}},
    }
    process_items(schema)
    assert schema["properties"]["nested"]["items"] == {"type": "object"}

    # Vertex rejects array types missing `items` entirely (not just empty).
    # Synthesize {"type": "object"} so the request validates.
    schema = {"type": "array"}
    process_items(schema)
    assert schema["items"] == {"type": "object"}

    # Test array missing items inside anyOf branch
    schema = {
        "type": "object",
        "properties": {
            "callbacks": {
                "anyOf": [{"type": "array"}, {"type": "object"}],
            }
        },
    }
    process_items(schema)
    assert schema["properties"]["callbacks"]["anyOf"][0]["items"] == {"type": "object"}


def test_build_vertex_schema_array_branch_missing_items_in_anyof():
    """
    Regression: an `anyOf` branch with `{"type": "array"}` (no items) must
    end up with synthesized `items: {"type": "object"}` after the schema
    transform — Vertex returns INVALID_ARGUMENT otherwise.
    """
    from litellm.llms.vertex_ai.common_utils import _build_vertex_schema

    parameters = {
        "properties": {
            "callbacks": {
                "anyOf": [
                    {"type": "array"},
                    {"type": "object"},
                    {"type": "null"},
                ]
            }
        },
        "type": "object",
    }

    result = _build_vertex_schema(parameters)
    callbacks_anyof = result["properties"]["callbacks"]["anyOf"]
    array_branches = [b for b in callbacks_anyof if b.get("type") == "array"]
    assert array_branches, "expected an array branch to remain after transform"
    for branch in array_branches:
        assert branch.get("items") == {
            "type": "object"
        }, f"array branch must have items synthesized; got {branch}"


def test_vertex_ai_complex_response_schema():
    import json
    from copy import deepcopy

    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
        VertexGeminiConfig,
    )

    v = VertexGeminiConfig()
    non_default_params = {
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "response_format": {
            "type": "json_schema",
            "json_schema": {
                "name": "generic_schema",
                "schema": {
                    "$id": "https://example.com/schema",
                    "type": "object",
                    "$schema": "https://json-schema.org/draft/2020-12/schema",
                    "required": ["field1", "field2", "field3"],
                    "properties": {
                        "field3": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "required": ["subfield1", "subfield2"],
                                "properties": {
                                    "subfield1": {
                                        "type": "string",
                                        "description": "string field",
                                    },
                                    "subfield2": {
                                        "type": "array",
                                        "items": {
                                            "type": "object",
                                            "oneOf": [
                                                {
                                                    "title": "Type1",
                                                    "required": ["Type1"],
                                                },
                                                {
                                                    "title": "Type2",
                                                    "required": ["Type2"],
                                                },
                                                {
                                                    "title": "Type3",
                                                    "required": ["Type3"],
                                                },
                                            ],
                                            "properties": {
                                                "Type1": {
                                                    "type": "object",
                                                    "required": ["prop1", "prop2"],
                                                    "properties": {
                                                        "prop1": {
                                                            "type": "string",
                                                            "description": "string value",
                                                        },
                                                        "prop2": {
                                                            "type": "string",
                                                            "description": "string value",
                                                        },
                                                    },
                                                    "additionalProperties": False,
                                                },
                                                "Type3": {
                                                    "type": "object",
                                                    "required": ["prop1", "prop3"],
                                                    "properties": {
                                                        "prop1": {
                                                            "type": "string",
                                                            "description": "string value",
                                                        },
                                                        "prop3": {
                                                            "type": "array",
                                                            "items": {
                                                                "type": "object",
                                                                "required": [
                                                                    "item1",
                                                                    "item2",
                                                                ],
                                                                "properties": {
                                                                    "item1": {
                                                                        "type": "string",
                                                                        "description": "string value",
                                                                    },
                                                                    "item2": {
                                                                        "type": "boolean",
                                                                        "description": "boolean value",
                                                                    },
                                                                },
                                                                "additionalProperties": False,
                                                            },
                                                        },
                                                    },
                                                    "additionalProperties": False,
                                                },
                                                "Type2": {
                                                    "type": "object",
                                                    "required": ["prop1", "prop3"],
                                                    "properties": {
                                                        "prop1": {
                                                            "type": "string",
                                                            "description": "string value",
                                                        },
                                                        "prop3": {
                                                            "type": "array",
                                                            "items": {
                                                                "type": "object",
                                                                "required": [
                                                                    "item1",
                                                                    "item2",
                                                                ],
                                                                "properties": {
                                                                    "item1": {
                                                                        "type": "string",
                                                                        "description": "string value",
                                                                    },
                                                                    "item2": {
                                                                        "type": "boolean",
                                                                        "description": "boolean value",
                                                                    },
                                                                },
                                                                "additionalProperties": False,
                                                            },
                                                        },
                                                    },
                                                    "additionalProperties": False,
                                                },
                                            },
                                            "additionalProperties": False,
                                        },
                                    },
                                },
                                "additionalProperties": False,
                            },
                        },
                        "field1": {"type": "string", "description": "string field"},
                        "field2": {"type": "string", "description": "string field"},
                    },
                    "additionalProperties": False,
                },
                "strict": True,
                "description": "Generic schema for testing",
            },
        },
    }
    original_non_default_params = deepcopy(non_default_params)
    optional_params = {}

    v.apply_response_schema_transformation(
        value=non_default_params["response_format"],
        optional_params=optional_params,
        model="gemini-1.5-pro-preview-0409",
    )

    # Assertions for the transformed schema
    transformed_schema = optional_params["response_schema"]

    # Check top level structure
    assert transformed_schema["type"] == "object"
    assert "propertyOrdering" in transformed_schema
    assert transformed_schema["propertyOrdering"] == ["field3", "field1", "field2"]
    assert set(transformed_schema["required"]) == {"field1", "field2", "field3"}

    # Check field3 structure (array of objects)
    field3 = transformed_schema["properties"]["field3"]
    assert field3["type"] == "array"

    # Check field3 items structure
    field3_items = field3["items"]
    assert field3_items["type"] == "object"
    assert "propertyOrdering" in field3_items
    assert field3_items["propertyOrdering"] == ["subfield1", "subfield2"]
    assert set(field3_items["required"]) == {"subfield1", "subfield2"}

    # Check subfield2 structure (array of objects)
    subfield2 = field3_items["properties"]["subfield2"]
    assert subfield2["type"] == "array"

    # Check subfield2 items structure
    subfield2_items = subfield2["items"]
    assert subfield2_items["type"] == "object"
    assert "propertyOrdering" in subfield2_items
    assert subfield2_items["propertyOrdering"] == ["Type1", "Type3", "Type2"]

    # Check Type1 structure
    type1 = subfield2_items["properties"]["Type1"]
    assert type1["type"] == "object"
    assert "propertyOrdering" in type1
    assert type1["propertyOrdering"] == ["prop1", "prop2"]
    assert set(type1["required"]) == {"prop1", "prop2"}

    # Check Type2 structure
    type2 = subfield2_items["properties"]["Type2"]
    assert type2["type"] == "object"
    assert "propertyOrdering" in type2
    assert type2["propertyOrdering"] == ["prop1", "prop3"]
    assert set(type2["required"]) == {"prop1", "prop3"}

    # Check Type3 structure
    type3 = subfield2_items["properties"]["Type3"]
    assert type3["type"] == "object"
    assert "propertyOrdering" in type3
    assert type3["propertyOrdering"] == ["prop1", "prop3"]
    assert set(type3["required"]) == {"prop1", "prop3"}

    # Check nested items in Type3's prop3
    type3_prop3 = type3["properties"]["prop3"]
    assert type3_prop3["type"] == "array"
    type3_prop3_items = type3_prop3["items"]
    assert type3_prop3_items["type"] == "object"
    assert "propertyOrdering" in type3_prop3_items
    assert type3_prop3_items["propertyOrdering"] == ["item1", "item2"]
    assert set(type3_prop3_items["required"]) == {"item1", "item2"}

    # Verify additionalProperties were removed during transformation
    assert "additionalProperties" not in transformed_schema
    assert "additionalProperties" not in field3_items
    assert "additionalProperties" not in subfield2_items
    assert "additionalProperties" not in type1
    assert "additionalProperties" not in type2
    assert "additionalProperties" not in type3
    assert "additionalProperties" not in type3_prop3_items


@pytest.mark.parametrize(
    "stream, expected_endpoint_suffix",
    [
        (True, "streamGenerateContent?alt=sse"),
        (False, "generateContent"),
    ],
)
def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
    """
    Test _get_vertex_url when vertex_location is 'global' for chat mode.
    """
    mode = "chat"
    model = "gemini-1.5-pro-preview-0409"
    vertex_project = "test-g-project"
    vertex_location = "global"
    vertex_api_version = "v1"

    # Mock litellm.VertexGeminiConfig.get_model_for_vertex_ai_url to return model as is
    # as we are not testing that part here, just the URL construction
    with patch(
        "litellm.VertexGeminiConfig.get_model_for_vertex_ai_url",
        side_effect=lambda model: model,
    ):
        url, endpoint = _get_vertex_url(
            mode=mode,
            model=model,
            stream=stream,
            vertex_project=vertex_project,
            vertex_location=vertex_location,
            vertex_api_version=vertex_api_version,
        )

    expected_url_base = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}"

    if stream:
        expected_endpoint = "streamGenerateContent"
        expected_url = f"{expected_url_base}:{expected_endpoint}?alt=sse"
    else:
        expected_endpoint = "generateContent"
        expected_url = f"{expected_url_base}:{expected_endpoint}"

    assert endpoint == expected_endpoint
    assert url == expected_url


@pytest.mark.parametrize(
    "model_cost_entry, vertex_region, expected_region",
    [
        # Model with supported_regions=["global"], no user region -> use "global"
        ({"supported_regions": ["global"]}, None, "global"),
        # Model with supported_regions=["global"], user passes unsupported region -> override to "global"
        ({"supported_regions": ["global"]}, "us-central1", "global"),
        # Model with supported_regions=["global"], user passes unsupported region -> override to "global"
        ({"supported_regions": ["global"]}, "europe-west1", "global"),
        # Model with supported_regions=["us-west2"], no user region -> use "us-west2"
        ({"supported_regions": ["us-west2"]}, None, "us-west2"),
        # Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it
        (
            {"supported_regions": ["us-west2", "us-central1"]},
            "us-central1",
            "us-central1",
        ),
        # Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override
        (
            {"supported_regions": ["us-west2", "us-central1"]},
            "europe-west1",
            "us-west2",
        ),
        # No model_cost entry, no user region -> default us-central1
        ({}, None, "us-central1"),
        # No model_cost entry, user specifies region -> use specified region
        ({}, "europe-west1", "europe-west1"),
        # No model_cost entry, user specifies region -> use specified region
        ({}, "us-east1", "us-east1"),
    ],
)
def test_get_vertex_region_global_only_model(
    model_cost_entry, vertex_region, expected_region
):
    """Test get_vertex_region resolves region from model_cost supported_regions"""
    import litellm
    from litellm.llms.vertex_ai.vertex_llm_base import VertexBase

    vertex_base = VertexBase()

    with patch.dict(
        litellm.model_cost,
        {"vertex_ai/test-model": model_cost_entry},
        clear=False,
    ):
        result = vertex_base.get_vertex_region(
            vertex_region=vertex_region, model="test-model"
        )

        assert result == expected_region


def test_vertex_filter_format_uri():
    import json

    from litellm.llms.vertex_ai.common_utils import filter_schema_fields

    parameters = {
        "type": "object",
        "properties": {
            "url": {
                "type": "string",
                "format": "uri",
                "description": "The URL to fetch content from",
            },
            "prompt": {
                "type": "string",
                "description": "The prompt to run on the fetched content",
            },
        },
        "required": ["url", "prompt"],
        "$schema": "http://json-schema.org/draft-07/schema#",
    }
    valid_schema_fields = {
        "minLength",
        "nullable",
        "maxItems",
        "required",
        "default",
        "items",
        "propertyOrdering",
        "maximum",
        "properties",
        "anyOf",
        "description",
        "minProperties",
        "minimum",
        "minItems",
        "maxProperties",
        "title",
        "pattern",
        "example",
        "format",
        "enum",
        "maxLength",
        "type",
    }

    new_parameters = filter_schema_fields(
        schema_dict=parameters,
        valid_fields=valid_schema_fields,
    )

    assert "uri" not in json.dumps(new_parameters)


def test_convert_schema_types_type_array_conversion():
    """
    Test _convert_schema_types function handles type arrays and case conversion.

    This test verifies the fix for the issue where type arrays like ["string", "number"]
    would raise an exception in Vertex AI schema validation.

    Relevant issue: https://github.com/BerriAI/litellm/issues/14091
    """
    from litellm.llms.vertex_ai.common_utils import _convert_schema_types

    # Input: OpenAI-style schema with type array (the problematic case)
    input_schema = {
        "type": "object",
        "properties": {
            "studio": {
                "type": ["string", "number"],
                "description": "The studio ID or name",
            }
        },
        "required": ["studio"],
        "additionalProperties": False,
        "$schema": "http://json-schema.org/draft-07/schema#",
    }

    # Expected output: Vertex AI compatible schema with anyOf and uppercase types
    expected_output = {
        "type": "object",
        "properties": {
            "studio": {
                "anyOf": [{"type": "string"}, {"type": "number"}],
                "description": "The studio ID or name",
            }
        },
        "required": ["studio"],
        "additionalProperties": False,
        "$schema": "http://json-schema.org/draft-07/schema#",
    }

    # Apply the transformation
    _convert_schema_types(input_schema)

    # Verify the transformation
    assert input_schema == expected_output

    # Verify specific transformations:
    # 1. Root level type converted to uppercase
    assert input_schema["type"] == "object"

    # 2. Type array converted to anyOf format
    assert "anyOf" in input_schema["properties"]["studio"]
    assert "type" not in input_schema["properties"]["studio"]

    # 3. Individual types in anyOf are uppercase
    anyof_types = input_schema["properties"]["studio"]["anyOf"]
    assert anyof_types[0]["type"] == "string"
    assert anyof_types[1]["type"] == "number"

    # 4. Other properties preserved
    assert (
        input_schema["properties"]["studio"]["description"] == "The studio ID or name"
    )
    assert input_schema["required"] == ["studio"]


def test_fix_enum_empty_strings():
    """
    Test _fix_enum_empty_strings function replaces empty strings with None in enum arrays.

    This test verifies the fix for the issue where Gemini rejects tool definitions
    with empty strings in enum values, causing API failures.

    Relevant issue: Gemini does not accept empty strings in enum values
    """
    from litellm.llms.vertex_ai.common_utils import _fix_enum_empty_strings

    # Input: Schema with empty string in enum (the problematic case)
    input_schema = {
        "type": "object",
        "properties": {
            "user_agent_type": {
                "enum": ["", "desktop", "mobile", "tablet"],
                "type": "string",
                "description": "Device type for user agent",
            }
        },
        "required": ["user_agent_type"],
    }

    # Expected output: Empty strings replaced with None
    expected_output = {
        "type": "object",
        "properties": {
            "user_agent_type": {
                "enum": [None, "desktop", "mobile", "tablet"],
                "type": "string",
                "description": "Device type for user agent",
            }
        },
        "required": ["user_agent_type"],
    }

    # Apply the transformation
    _fix_enum_empty_strings(input_schema)

    # Verify the transformation
    assert input_schema == expected_output

    # Verify specific transformations:
    # 1. Empty string replaced with None
    enum_values = input_schema["properties"]["user_agent_type"]["enum"]
    assert "" not in enum_values
    assert None in enum_values

    # 2. Other enum values preserved
    assert "desktop" in enum_values
    assert "mobile" in enum_values
    assert "tablet" in enum_values


def test_get_vertex_model_id_from_url():
    """Test get_vertex_model_id_from_url with various URLs"""
    from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url

    # Test with valid URL
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id == "gemini-pro"

    # Test with invalid URL
    url = "https://invalid-url.com"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id is None


def test_get_vertex_model_id_from_url_with_slashes():
    """Test get_vertex_model_id_from_url with model names containing slashes (e.g., gcp/google/gemini-2.5-flash)

    Regression test for NVIDIA issue: custom model names with slashes in passthrough URLs
    were being truncated (e.g., 'gcp/google/gemini-2.5-flash' -> 'gcp'), causing access_groups
    checks to fail.
    """
    from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url

    # Test with model name containing slashes: gcp/google/gemini-2.5-flash
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gcp/google/gemini-2.5-flash:generateContent"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id == "gcp/google/gemini-2.5-flash"

    # Test with model name containing slashes: gcp/google/gemini-3-flash-preview
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/google/models/gcp/google/gemini-3-flash-preview:streamGenerateContent"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id == "gcp/google/gemini-3-flash-preview"

    # Test with custom model path: custom/model
    url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/custom/model:generateContent"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id == "custom/model"

    # Test passthrough URL format (without host)
    url = "v1/projects/my-project/locations/us-central1/publishers/google/models/gcp/google/gemini-2.5-flash:generateContent"
    model_id = get_vertex_model_id_from_url(url)
    assert model_id == "gcp/google/gemini-2.5-flash"


def test_construct_target_url_with_version_prefix():
    """Test construct_target_url with version prefixes"""
    from litellm.llms.vertex_ai.common_utils import construct_target_url

    # Test with /v1/ prefix
    url = "/v1/publishers/google/models/gemini-pro:streamGenerateContent"
    vertex_project = "test-project"
    vertex_location = "us-central1"
    base_url = "https://us-central1-aiplatform.googleapis.com"

    target_url = construct_target_url(
        base_url=base_url,
        requested_route=url,
        vertex_project=vertex_project,
        vertex_location=vertex_location,
    )

    expected_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
    assert str(target_url) == expected_url

    # Test with /v1beta1/ prefix
    url = "/v1beta1/publishers/google/models/gemini-pro:streamGenerateContent"

    target_url = construct_target_url(
        base_url=base_url,
        requested_route=url,
        vertex_project=vertex_project,
        vertex_location=vertex_location,
    )

    expected_url = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:streamGenerateContent"
    assert str(target_url) == expected_url


def test_fix_enum_types():
    """
    Test _fix_enum_types function removes enum fields when type is not string.

    This test verifies the fix for the issue where Gemini rejects cached content
    with function parameter enums on non-string types, causing API failures.

    Relevant issue: Gemini only allows enums for string-typed fields
    """
    from litellm.llms.vertex_ai.common_utils import _fix_enum_types

    # Input: Schema with enum on non-string type (the problematic case)
    input_schema = {
        "type": "object",
        "properties": {
            "truncateMode": {
                "enum": ["auto", "none", "start", "end"],
                "type": "string",  # This should keep the enum
                "description": "How to truncate content",
            },
            "maxLength": {
                "enum": [100, 200, 500],  # This should be removed
                "type": "integer",
                "description": "Maximum length",
            },
            "enabled": {
                "enum": [True, False],  # This should be removed
                "type": "boolean",
                "description": "Whether feature is enabled",
            },
            "nested": {
                "type": "object",
                "properties": {
                    "innerEnum": {
                        "enum": ["a", "b", "c"],  # This should be kept
                        "type": "string",
                    },
                    "innerNonStringEnum": {
                        "enum": [1, 2, 3],  # This should be removed
                        "type": "integer",
                    },
                },
            },
            "anyOfField": {
                "anyOf": [
                    {
                        "type": "string",
                        "enum": ["option1", "option2"],
                    },  # This should be kept
                    {"type": "integer", "enum": [1, 2, 3]},  # This should be removed
                ]
            },
        },
    }

    # Expected output: Non-string enums removed, string enums kept
    expected_output = {
        "type": "object",
        "properties": {
            "truncateMode": {
                "enum": ["auto", "none", "start", "end"],  # Kept - string type
                "type": "string",
                "description": "How to truncate content",
            },
            "maxLength": {  # enum removed
                "type": "integer",
                "description": "Maximum length",
            },
            "enabled": {  # enum removed
                "type": "boolean",
                "description": "Whether feature is enabled",
            },
            "nested": {
                "type": "object",
                "properties": {
                    "innerEnum": {
                        "enum": ["a", "b", "c"],  # Kept - string type
                        "type": "string",
                    },
                    "innerNonStringEnum": {"type": "integer"},  # enum removed
                },
            },
            "anyOfField": {
                "anyOf": [
                    {
                        "type": "string",
                        "enum": ["option1", "option2"],
                    },  # Kept - has string type
                    {"type": "integer"},  # enum removed
                ]
            },
        },
    }

    # Apply the transformation
    _fix_enum_types(input_schema)

    # Verify the transformation
    assert input_schema == expected_output

    # Verify specific transformations:
    # 1. String enums are preserved
    assert "enum" in input_schema["properties"]["truncateMode"]
    assert input_schema["properties"]["truncateMode"]["enum"] == [
        "auto",
        "none",
        "start",
        "end",
    ]

    assert "enum" in input_schema["properties"]["nested"]["properties"]["innerEnum"]
    assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == [
        "a",
        "b",
        "c",
    ]

    # 2. Non-string enums are removed
    assert "enum" not in input_schema["properties"]["maxLength"]
    assert "enum" not in input_schema["properties"]["enabled"]
    assert (
        "enum"
        not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
    )

    # 3. anyOf with string type keeps enum, non-string removes it
    assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0]
    assert "enum" not in input_schema["properties"]["anyOfField"]["anyOf"][1]

    # 4. Other properties preserved
    assert input_schema["properties"]["maxLength"]["type"] == "integer"
    assert input_schema["properties"]["enabled"]["type"] == "boolean"


def test_get_token_url():
    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
        VertexLLM,
    )

    vertex_llm = VertexLLM()
    vertex_ai_project = "pathrise-convert-1606954137718"
    vertex_ai_location = "us-central1"
    vertex_credentials = ""

    should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
        optional_params={"cached_content": "hi"}
    )

    _, url = vertex_llm._get_token_and_url(
        auth_header=None,
        vertex_project=vertex_ai_project,
        vertex_location=vertex_ai_location,
        vertex_credentials=vertex_credentials,
        gemini_api_key="",
        custom_llm_provider="vertex_ai_beta",
        should_use_v1beta1_features=should_use_v1beta1_features,
        api_base=None,
        model="",
        stream=False,
    )

    print("url=", url)

    should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
        optional_params={"temperature": 0.1}
    )

    _, url = vertex_llm._get_token_and_url(
        auth_header=None,
        vertex_project=vertex_ai_project,
        vertex_location=vertex_ai_location,
        vertex_credentials=vertex_credentials,
        gemini_api_key="",
        custom_llm_provider="vertex_ai_beta",
        should_use_v1beta1_features=should_use_v1beta1_features,
        api_base=None,
        model="",
        stream=False,
    )

    print("url for normal request", url)

    assert "v1beta1" not in url
    assert "/v1/" in url

    pass


@pytest.mark.asyncio
async def test_vertex_ai_token_counter_routes_partner_models():
    """
    Test that VertexAITokenCounter correctly routes partner models (Claude, Mistral, etc.)
    to the partner models token counter instead of the Gemini token counter.
    """
    from unittest.mock import AsyncMock, patch

    from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter
    from litellm.types.utils import TokenCountResponse

    token_counter = VertexAITokenCounter()

    # Mock the partner models handler
    with patch(
        "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels.count_tokens"
    ) as mock_partner_count_tokens:
        mock_partner_count_tokens.return_value = {
            "input_tokens": 42,
            "tokenizer_used": "vertex_ai_partner_models",
        }

        # Test with a Claude model (partner model)
        result = await token_counter.count_tokens(
            model_to_use="claude-3-5-sonnet-20241022",
            messages=[{"role": "user", "content": "Hello"}],
            contents=None,
            deployment={
                "litellm_params": {
                    "vertex_project": "test-project",
                    "vertex_location": "us-east5",
                }
            },
            request_model="vertex_ai/claude-3-5-sonnet-20241022",
        )

        # Verify partner models handler was called
        assert mock_partner_count_tokens.called
        assert result is not None
        assert isinstance(result, TokenCountResponse)
        assert result.total_tokens == 42
        assert result.tokenizer_type == "vertex_ai_partner_models"


@pytest.mark.asyncio
async def test_vertex_ai_token_counter_uses_count_tokens_location():
    """
    Test that VertexAITokenCounter uses vertex_count_tokens_location to override
    vertex_location when counting tokens for partner models.

    Count tokens API is not available on global location for partner models:
    https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens
    """
    from unittest.mock import patch

    from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter
    from litellm.types.utils import TokenCountResponse

    token_counter = VertexAITokenCounter()

    # Mock the partner models handler
    with patch(
        "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels.count_tokens"
    ) as mock_partner_count_tokens:
        mock_partner_count_tokens.return_value = {
            "input_tokens": 42,
            "tokenizer_used": "vertex_ai_partner_models",
        }

        # Test with vertex_count_tokens_location overriding vertex_location
        await token_counter.count_tokens(
            model_to_use="claude-3-5-sonnet-20241022",
            messages=[{"role": "user", "content": "Hello"}],
            contents=None,
            deployment={
                "litellm_params": {
                    "vertex_project": "test-project",
                    "vertex_location": "global",  # Original location (not supported for count_tokens)
                    "vertex_count_tokens_location": "us-east5",  # Override for count_tokens
                }
            },
            request_model="vertex_ai/claude-3-5-sonnet-20241022",
        )

        # Verify the partner models handler was called with the overridden location
        assert mock_partner_count_tokens.called
        call_kwargs = mock_partner_count_tokens.call_args.kwargs
        assert call_kwargs["vertex_location"] == "us-east5"
        assert call_kwargs["vertex_project"] == "test-project"


@pytest.mark.asyncio
async def test_vertex_ai_token_counter_routes_gemini_models():
    """
    Test that VertexAITokenCounter correctly routes Gemini models
    to the Gemini token counter (not partner models).
    """
    from unittest.mock import AsyncMock, patch

    from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter
    from litellm.types.utils import TokenCountResponse

    token_counter = VertexAITokenCounter()

    # Mock the Gemini handler (different import path)
    with patch(
        "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens"
    ) as mock_gemini_count_tokens:
        mock_gemini_count_tokens.return_value = {
            "totalTokens": 50,
            "tokenizer_used": "gemini",
        }

        # Test with a Gemini model (not a partner model)
        result = await token_counter.count_tokens(
            model_to_use="gemini-1.5-pro",
            messages=[{"role": "user", "content": "Hello"}],
            contents=None,
            deployment={
                "litellm_params": {
                    "vertex_project": "test-project",
                    "vertex_location": "us-central1",
                }
            },
            request_model="vertex_ai/gemini-1.5-pro",
        )

        # Verify Gemini handler was called
        assert mock_gemini_count_tokens.called
        assert result is not None
        assert isinstance(result, TokenCountResponse)
        assert result.total_tokens == 50


@pytest.mark.asyncio
async def test_vertex_ai_partner_model_detection():
    """
    Test that VertexAIPartnerModels.is_vertex_partner_model correctly identifies
    partner models (Claude, Mistral, Llama, etc.).
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    # Test Claude models (should be detected as partner model)
    assert VertexAIPartnerModels.is_vertex_partner_model("claude-3-5-sonnet-20241022")
    assert VertexAIPartnerModels.is_vertex_partner_model("claude-3-opus-20240229")
    assert VertexAIPartnerModels.is_vertex_partner_model("claude-3-haiku-20240307")

    # Test Mistral models
    assert VertexAIPartnerModels.is_vertex_partner_model("mistral-large-2407")
    assert VertexAIPartnerModels.is_vertex_partner_model("mistral-7b-instruct-v0.3")

    # Test Meta/Llama models
    assert VertexAIPartnerModels.is_vertex_partner_model("meta/llama-3.1-405b")
    # Test Minimax models
    assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas")
    # Test Moonshot models
    assert VertexAIPartnerModels.is_vertex_partner_model(
        "moonshotai/kimi-k2-thinking-maas"
    )

    # Test Gemini models (should NOT be detected as partner model)
    assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro")
    assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.0-pro")
    assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-pro-vision")

    # Test other non-partner models
    assert not VertexAIPartnerModels.is_vertex_partner_model("text-bison-001")
    assert not VertexAIPartnerModels.is_vertex_partner_model("chat-bison-001")


def test_vertex_ai_minimax_uses_openai_handler():
    """
    Ensure Minimax partner models re-use the OpenAI-format handler.
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.should_use_openai_handler("minimaxai/minimax-m2-maas")


def test_vertex_ai_moonshot_uses_openai_handler():
    """
    Ensure Moonshot partner models re-use the OpenAI-format handler.
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.should_use_openai_handler(
        "moonshotai/kimi-k2-thinking-maas"
    )


def test_vertex_ai_zai_uses_openai_handler():
    """
    Ensure ZAI partner models re-use the OpenAI-format handler.
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.should_use_openai_handler("zai-org/glm-4.7-maas")


def test_vertex_ai_zai_is_partner_model():
    """
    Ensure ZAI models are detected as Vertex AI partner models.
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.is_vertex_partner_model("zai-org/glm-4.7-maas")


def test_vertex_ai_gemma_maas_is_partner_model():
    """
    Ensure Gemma MaaS models are detected as Vertex AI partner models so they
    route through the OpenAI-compatible /endpoints/openapi path (not the
    legacy non-gemini path or the vertex_ai/gemma/ predict-endpoint handler).
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.is_vertex_partner_model(
        "google/gemma-4-26b-a4b-it-maas"
    )


def test_vertex_ai_gemma_maas_uses_openai_handler():
    """
    Ensure Gemma MaaS partner models re-use the OpenAI-format handler.
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert VertexAIPartnerModels.should_use_openai_handler(
        "google/gemma-4-26b-a4b-it-maas"
    )


def test_vertex_ai_gemma_maas_routes_to_partner_models():
    """
    Regression guard for owtaylor's worry that Gemma MaaS could be misrouted as
    a gemma model. get_vertex_ai_model_route must return PARTNER_MODELS, never
    GEMMA, MODEL_GARDEN, or NON_GEMINI.
    """
    from litellm.llms.vertex_ai.common_utils import (
        VertexAIModelRoute,
        get_vertex_ai_model_route,
    )

    route = get_vertex_ai_model_route("google/gemma-4-26b-a4b-it-maas")
    assert route == VertexAIModelRoute.PARTNER_MODELS


def test_vertex_ai_google_gemini_not_detected_as_gemma_maas():
    """
    Negative: adding the "google/gemma-" prefix must not widen detection to
    other google/* models like google/gemini-* (which should keep flowing
    through the gemini route, not partner_models).
    """
    from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
        VertexAIPartnerModels,
    )

    assert not VertexAIPartnerModels.is_vertex_partner_model("google/gemini-1.5-pro")
    assert not VertexAIPartnerModels.should_use_openai_handler("google/gemini-1.5-pro")


def test_build_vertex_schema_empty_properties():
    """
    Test _build_vertex_schema handles empty properties objects correctly.

    This test verifies the fix for the issue where Gemini rejects schemas
    with empty properties objects like {"properties": {}, "type": "object"}.

    Error from Gemini: "GenerateContentRequest.generation_config.response_schema
    .properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties:
    should be non-empty for OBJECT type"

    The fix removes empty properties objects and their associated type/required fields.
    """
    from litellm.llms.vertex_ai.common_utils import _build_vertex_schema

    # Input: Schema with empty properties (the problematic case from real request)
    input_schema = {
        "properties": {
            "action": {
                "description": "List of actions to execute",
                "items": {
                    "anyOf": [
                        {
                            "properties": {
                                "go_back": {
                                    "properties": {},
                                    "type": "object",
                                    "additionalProperties": False,
                                    "description": "Go back",
                                    "required": [],
                                }
                            },
                            "required": ["go_back"],
                            "type": "object",
                            "additionalProperties": False,
                        }
                    ]
                },
                "type": "array",
            }
        },
        "type": "object",
        "additionalProperties": False,
    }

    # Apply the transformation
    result = _build_vertex_schema(input_schema)

    # Verify the transformation removed empty properties
    # Navigate to the go_back schema
    go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"][
        "go_back"
    ]

    # Verify empty properties was removed
    assert "properties" not in go_back_schema, "Empty properties should be removed"

    # Verify type is kept as object (Gemini requires type: object even without properties)
    assert (
        go_back_schema.get("type") == "object"
    ), "Type should be kept as object when properties is empty"

    # Verify required was also removed
    assert (
        "required" not in go_back_schema
    ), "Required should be removed when properties is empty"

    # Verify description is preserved
    assert (
        go_back_schema.get("description") == "Go back"
    ), "Description should be preserved"

    # Verify parent schema still has proper structure
    parent_schema = result["properties"]["action"]["items"]["anyOf"][0]
    assert (
        parent_schema["type"] == "object"
    ), "Parent schema should still have object type"
    assert (
        "go_back" in parent_schema["properties"]
    ), "go_back should still be in parent properties"


def test_add_object_type_schema_with_no_properties_and_no_type():
    """
    Test that add_object_type adds type: object when schema has no properties and no type.
    Fixes issue where tools with no arguments (e.g. EnterPlanMode) fail on Gemini.
    """
    from litellm.llms.vertex_ai.common_utils import add_object_type

    # Input: Schema with no properties and no type (the problematic case)
    input_schema = {"$schema": "https://json-schema.org/draft/2020-12/schema"}

    # Apply the transformation
    add_object_type(input_schema)

    # Verify type: object was added
    assert input_schema.get("type") == "object", "type: object should be added"

    # Verify $schema is preserved
    assert input_schema.get("$schema") == "https://json-schema.org/draft/2020-12/schema"


def test_add_object_type_does_not_override_existing_type():
    """
    Test add_object_type does not override existing type field.
    """
    from litellm.llms.vertex_ai.common_utils import add_object_type

    # Input: Schema with existing type
    input_schema = {"type": "string", "description": "A string field"}

    # Apply the transformation
    add_object_type(input_schema)

    # Verify type was not changed
    assert input_schema.get("type") == "string", "Existing type should not be changed"


def test_add_object_type_does_not_add_type_when_anyof_present():
    """
    Test add_object_type does not add type: object when anyOf is present.
    """
    from litellm.llms.vertex_ai.common_utils import add_object_type

    # Input: Schema with anyOf but no type
    input_schema = {"anyOf": [{"type": "string"}, {"type": "null"}]}

    # Apply the transformation
    add_object_type(input_schema)

    # Verify type was not added (anyOf handles the type)
    assert "type" not in input_schema, "type should not be added when anyOf is present"


def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata():
    assert vertex_request_labels_from_litellm_params(None) is None
    assert vertex_request_labels_from_litellm_params({}) is None
    assert vertex_request_labels_from_litellm_params({"metadata": None}) is None
    lp = {"metadata": {"requester_metadata": {"team": "analytics", "count": 3}}}
    assert vertex_request_labels_from_litellm_params(lp) == {"team": "analytics"}


def test_vertex_request_labels_from_litellm_params_accepts_litellm_metadata():
    lp = {
        "litellm_metadata": {
            "requester_metadata": {"team": "platform", "count": 3}
        }
    }
    assert vertex_request_labels_from_litellm_params(lp) == {"team": "platform"}


def test_vertex_request_labels_prefers_metadata_over_litellm_metadata():
    lp = {
        "metadata": {"requester_metadata": {"source": "metadata"}},
        "litellm_metadata": {"requester_metadata": {"source": "litellm_metadata"}},
    }
    assert vertex_request_labels_from_litellm_params(lp) == {"source": "metadata"}


def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
    optional = {"labels": {"env": "prod"}}
    litellm_params = {"metadata": {"requester_metadata": {"team": "x"}}}
    assert pop_vertex_request_labels(optional, litellm_params) == {"env": "prod"}
    assert "labels" not in optional

    optional2: dict = {}
    assert pop_vertex_request_labels(optional2, litellm_params) == {"team": "x"}

    optional3 = {"labels": {"team": 123}}
    assert pop_vertex_request_labels(optional3, litellm_params) == {"team": "x"}


def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent():
    optional: dict = {}
    litellm_params = {
        "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}
    }
    assert pop_vertex_request_labels(optional, litellm_params) == {
        "team": "from_litellm_meta"
    }


def test_vertex_text_embedding_request_includes_labels_from_metadata():
    import litellm

    req = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
        input="hi",
        optional_params={},
        model="text-embedding-004",
        litellm_params={
            "metadata": {"requester_metadata": {"project_id": "cost-center-1"}}
        },
    )
    assert req.get("labels") == {"project_id": "cost-center-1"}
