"""
Tests for Vertex AI Qwen MaaS models that require the global endpoint.

These tests verify that:
1. The correct global URL is constructed (https://aiplatform.googleapis.com)
2. The get_vertex_region method resolves regions from model_cost supported_regions
3. The completion() and responses() API work with Qwen models
"""

import os
from unittest.mock import MagicMock, patch, AsyncMock

import pytest


import litellm
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VertexPartnerProvider


@pytest.fixture(autouse=True)
def clean_vertex_env():
    """Clear Google/Vertex AI environment variables before each test to prevent test isolation issues."""
    saved_env = {}
    env_vars_to_clear = [
        "GOOGLE_APPLICATION_CREDENTIALS",
        "GOOGLE_CLOUD_PROJECT",
        "VERTEXAI_PROJECT",
        "VERTEX_PROJECT",
        "VERTEX_LOCATION",
        "VERTEX_AI_PROJECT",
    ]
    for var in env_vars_to_clear:
        if var in os.environ:
            saved_env[var] = os.environ[var]
            del os.environ[var]

    yield

    # Restore saved environment variables
    for var, value in saved_env.items():
        os.environ[var] = value


class TestVertexBaseGetVertexRegion:
    """Test the get_vertex_region method using model_cost lookup."""

    def test_global_model_no_user_region_returns_global(self):
        """Test that global-only models return 'global' when user doesn't specify region."""
        vertex_base = VertexBase()

        with patch.dict(
            litellm.model_cost,
            {
                "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
                    "supported_regions": ["global"]
                }
            },
            clear=False,
        ):
            result = vertex_base.get_vertex_region(
                vertex_region=None,
                model="qwen/qwen3-next-80b-a3b-instruct-maas",
            )
            assert result == "global"

    def test_global_model_with_unsupported_user_region_overrides(self):
        """Test that unsupported user region is overridden for global-only models."""
        vertex_base = VertexBase()

        with patch.dict(
            litellm.model_cost,
            {
                "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
                    "supported_regions": ["global"]
                }
            },
            clear=False,
        ):
            result = vertex_base.get_vertex_region(
                vertex_region="us-central1",
                model="qwen/qwen3-next-80b-a3b-instruct-maas",
            )
            assert result == "global"

    def test_non_global_model_uses_provided_region(self):
        """Test that non-global models use the provided region."""
        vertex_base = VertexBase()

        with patch.dict(litellm.model_cost, {}, clear=False):
            result = vertex_base.get_vertex_region(
                vertex_region="europe-west1",
                model="gemini-1.5-pro",
            )
            assert result == "europe-west1"

    def test_non_global_model_fallback_to_us_central1(self):
        """Test that non-global models with None region fallback to us-central1."""
        vertex_base = VertexBase()

        with patch.dict(litellm.model_cost, {}, clear=False):
            result = vertex_base.get_vertex_region(
                vertex_region=None,
                model="unknown-model-xyz",
            )
            assert result == "us-central1"


class TestCreateVertexURLGlobal:
    """Test that create_vertex_url handles global location correctly."""

    def test_global_location_url_format(self):
        """Test that global location produces correct URL without region prefix."""
        url = VertexBase.create_vertex_url(
            vertex_location="global",
            vertex_project="test-project",
            partner=VertexPartnerProvider.llama,
            stream=False,
            model="qwen/qwen3-next-80b-a3b-instruct-maas",
        )

        # Global URL should NOT have region prefix
        assert url.startswith("https://aiplatform.googleapis.com")
        assert "global-aiplatform.googleapis.com" not in url
        assert "/locations/global/" in url

    def test_regional_location_url_format(self):
        """Test that regional location produces correct URL with region prefix."""
        url = VertexBase.create_vertex_url(
            vertex_location="us-central1",
            vertex_project="test-project",
            partner=VertexPartnerProvider.llama,
            stream=False,
            model="openai/gpt-oss-20b-maas",
        )

        # Regional URL should have region prefix
        assert url.startswith("https://us-central1-aiplatform.googleapis.com")
        assert "/locations/us-central1/" in url


@pytest.mark.asyncio
async def test_vertex_ai_qwen_global_endpoint_url():
    """
    Test that Qwen models use the global endpoint URL.
    """
    # Mock response
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.headers = {}
    mock_response.json.return_value = {
        "id": "chatcmpl-qwen-test",
        "object": "chat.completion",
        "created": 1234567890,
        "model": "qwen/qwen3-next-80b-a3b-instruct-maas",
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": "Hello! How can I help you today?",
                },
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
    }

    mock_vertexai = MagicMock()
    mock_vertexai.preview = MagicMock()

    with (
        patch(
            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler"
        ) as mock_http_handler,
        patch(
            "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token",
            return_value=("fake-token", "test-project"),
        ),
        patch.dict(
            "sys.modules",
            {"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview},
        ),
        patch.dict(
            litellm.model_cost,
            {
                "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
                    "supported_regions": ["global"]
                }
            },
            clear=False,
        ),
    ):
        mock_http_handler.return_value.post = AsyncMock(return_value=mock_response)

        response = await litellm.acompletion(
            model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
            messages=[{"role": "user", "content": "Hello"}],
            vertex_ai_project="test-project",
        )

        # Verify the mock was called
        mock_http_handler.return_value.post.assert_called_once()

        # Get the call arguments
        call_args = mock_http_handler.return_value.post.call_args
        called_url = call_args.kwargs["url"]

        # Verify the URL uses global endpoint (no region prefix)
        assert called_url.startswith("https://aiplatform.googleapis.com")
        assert "global-aiplatform.googleapis.com" not in called_url
        assert "/locations/global/" in called_url
        assert "/endpoints/openapi/chat/completions" in called_url

        # Verify response
        assert response.model == "qwen/qwen3-next-80b-a3b-instruct-maas"


class TestGetSupportedRegions:
    """Test that get_supported_regions correctly reads from model_cost."""

    def test_get_supported_regions_returns_list(self):
        """Test that get_supported_regions returns a list when model has supported_regions."""
        # Mock the model_cost to have supported_regions
        with patch.dict(
            litellm.model_cost,
            {
                "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {
                    "supported_regions": ["global"],
                    "litellm_provider": "vertex_ai-qwen_models",
                }
            },
        ):
            regions = litellm.utils.get_supported_regions(
                model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
                custom_llm_provider="vertex_ai",
            )
            assert regions == ["global"]

    def test_get_supported_regions_returns_none_when_not_set(self):
        """Test that get_supported_regions returns None when model doesn't have supported_regions."""
        # Mock the model_cost without supported_regions
        with patch.dict(
            litellm.model_cost,
            {
                "vertex_ai/gemini-1.5-pro": {
                    "litellm_provider": "vertex_ai",
                }
            },
        ):
            regions = litellm.utils.get_supported_regions(
                model="vertex_ai/gemini-1.5-pro",
                custom_llm_provider="vertex_ai",
            )
            assert regions is None

    def test_get_supported_regions_returns_none_for_unknown_model(self):
        """Test that get_supported_regions returns None for unknown models."""
        regions = litellm.utils.get_supported_regions(
            model="vertex_ai/unknown-model-xyz",
            custom_llm_provider="vertex_ai",
        )
        assert regions is None
