"""
Unit tests for VertexAIRealtimeConfig.

Validates:
- URL construction (regional and global)
- Auth headers (Bearer token + project header)
- Session setup message format
- Full text-in / text-out round-trip via RealTimeStreaming with a mocked
  WebSocket pair (no real network calls)
"""

import json
from unittest.mock import AsyncMock, MagicMock

import pytest
import websockets.exceptions  # registers websockets.exceptions on the websockets namespace

import litellm
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig

# ---------------------------------------------------------------------------
# Config unit tests
# ---------------------------------------------------------------------------


def test_get_complete_url_regional():
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    url = cfg.get_complete_url(api_base=None, model="gemini-2.0-flash-live-001")
    assert url == (
        "wss://us-central1-aiplatform.googleapis.com"
        "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
    )


def test_get_complete_url_global():
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="global"
    )
    url = cfg.get_complete_url(api_base=None, model="gemini-2.0-flash-live-001")
    assert url == (
        "wss://aiplatform.googleapis.com"
        "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
    )


def test_get_complete_url_custom_api_base():
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    url = cfg.get_complete_url(
        api_base="https://custom-gateway.example.com",
        model="gemini-2.0-flash-live-001",
    )
    assert url.startswith("wss://custom-gateway.example.com")
    assert "BidiGenerateContent" in url


def test_validate_environment_sets_bearer_and_project():
    cfg = VertexAIRealtimeConfig(
        access_token="mytoken", project="proj-123", location="us-central1"
    )
    headers = cfg.validate_environment(
        headers={}, model="gemini-2.0-flash-live-001", api_key=None
    )
    assert headers["Authorization"] == "Bearer mytoken"
    assert headers["x-goog-user-project"] == "proj-123"


def test_session_configuration_request_model_format():
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    raw = cfg.session_configuration_request("gemini-2.0-flash-live-001")
    parsed = json.loads(raw)
    assert parsed["setup"]["model"] == (
        "projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-001"
    )


def test_vertex_requires_session_configuration_feature_flag(monkeypatch):
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    # Default remains backwards-compatible (auto setup on connect)
    monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
    assert cfg.requires_session_configuration() is True

    # Opt-in deferred setup for tool-injection flow
    monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
    assert cfg.requires_session_configuration() is False


def test_vertex_session_update_defaults_to_audio_modality():
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    session_update = {
        "type": "session.update",
        "session": {
            "instructions": "You are a helpful assistant.",
            # No modalities provided on purpose
        },
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        "gemini-live-2.5-flash-preview-native-audio-09-2025",
        session_configuration_request=None,
    )
    assert len(messages) == 1
    setup_payload = json.loads(messages[0])["setup"]
    assert setup_payload["generationConfig"]["responseModalities"] == ["AUDIO"]


_NATIVE_AUDIO_MODEL = "gemini-live-2.5-flash-preview-native-audio-09-2025"


@pytest.fixture(autouse=False)
def patch_native_audio_cost_map_entry(monkeypatch):
    """Inject gemini_native_audio into the cost map for the test model.

    litellm.model_cost is fetched from main branch at import time, so in CI
    the field may not exist yet. Patch it locally so these unit tests remain
    self-contained and don't depend on the remote cost map state.
    """
    entry = dict(litellm.model_cost.get(_NATIVE_AUDIO_MODEL, {}))
    entry["gemini_native_audio"] = True
    monkeypatch.setitem(litellm.model_cost, _NATIVE_AUDIO_MODEL, entry)


def test_vertex_audio_only_live_model_coerces_text_modality_to_audio(
    patch_native_audio_cost_map_entry,
):
    """Regression: TEXT-only responseModalities causes 1007 on native-audio Live models."""
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    session_update = {
        "type": "session.update",
        "session": {
            "modalities": ["text"],
            "instructions": "You are a terse assistant.",
        },
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )

    setup = json.loads(messages[0])["setup"]
    assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]


def test_vertex_session_update_normalizes_ga_remapped_fields(
    patch_native_audio_cost_map_entry,
):
    """GA-format clients send ``output_modalities`` and nested
    ``audio.input.transcription`` / ``audio.input.turn_detection``. These must
    be normalised back to the flat beta keys before ``map_openai_params``
    runs so client preferences aren't silently dropped.
    """
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    session_update = {
        "type": "session.update",
        "session": {
            "instructions": "Be concise.",
            "output_modalities": ["text"],
            "audio": {
                "input": {
                    "transcription": {},
                    "turn_detection": {"silence_duration_ms": 1500},
                },
            },
        },
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )
    assert len(messages) == 1
    setup_payload = json.loads(messages[0])["setup"]

    assert setup_payload["generationConfig"]["responseModalities"] == ["AUDIO"]
    assert setup_payload["inputAudioTranscription"] == {}
    assert (
        setup_payload["realtimeInputConfig"]["automaticActivityDetection"][
            "silenceDurationMs"
        ]
        == 1500
    )


# ---------------------------------------------------------------------------
# Round-trip test: text-in / text-out via RealTimeStreaming
# ---------------------------------------------------------------------------

# Minimal Gemini BidiGenerateContent message sequence:
#   server → setupComplete
#   client → conversation.item.create  (OpenAI format, translated by config)
#   server → serverContent with modelTurn text delta
#   server → serverContent with generationComplete

SETUP_COMPLETE = json.dumps({"setupComplete": {}})

SERVER_TEXT_DELTA = json.dumps(
    {"serverContent": {"modelTurn": {"parts": [{"text": "Hello from Vertex AI!"}]}}}
)

# generationComplete fires RESPONSE_TEXT_DONE; turnComplete fires RESPONSE_DONE
# They must be separate messages (the transformer processes one top-level key per message).
SERVER_GENERATION_COMPLETE = json.dumps({"serverContent": {"generationComplete": True}})

SERVER_TURN_COMPLETE = json.dumps({"serverContent": {"turnComplete": True}})

# OpenAI-format text message the client sends
CLIENT_TEXT_MESSAGE = json.dumps(
    {
        "type": "conversation.item.create",
        "item": {
            "type": "message",
            "role": "user",
            "content": [{"type": "input_text", "text": "Say hello"}],
        },
    }
)


@pytest.mark.asyncio
async def test_vertex_realtime_text_in_text_out():
    """
    Simulate a full text-in / text-out session through RealTimeStreaming using
    VertexAIRealtimeConfig for message translation.  All I/O is mocked.
    """
    from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming

    cfg = VertexAIRealtimeConfig(
        access_token="fake-token",
        project="fake-project",
        location="us-central1",
    )

    # --- mock client WebSocket (FastAPI side) ---
    client_ws = MagicMock()
    client_ws.exceptions = MagicMock()
    client_ws.exceptions.ConnectionClosed = Exception

    sent_to_client: list[str] = []

    async def _client_send_text(data: str):
        sent_to_client.append(data)

    client_ws.send_text = AsyncMock(side_effect=_client_send_text)

    # Client sends one text message then raises to end the loop
    client_ws.receive_text = AsyncMock(
        side_effect=[CLIENT_TEXT_MESSAGE, Exception("client done")]
    )

    # --- mock backend WebSocket (Vertex AI side) ---
    backend_ws = MagicMock()

    upstream_messages = [
        SETUP_COMPLETE,
        SERVER_TEXT_DELTA,
        SERVER_GENERATION_COMPLETE,
        SERVER_TURN_COMPLETE,
    ]

    async def _backend_recv(decode=True):
        if not upstream_messages:
            # Signal normal connection close so the loop exits cleanly
            raise websockets.exceptions.ConnectionClosedOK(None, None)  # type: ignore[arg-type]
        return upstream_messages.pop(0)

    backend_ws.recv = AsyncMock(side_effect=_backend_recv)

    sent_to_backend: list[str] = []

    async def _backend_send(data: str):
        sent_to_backend.append(data)

    backend_ws.send = AsyncMock(side_effect=_backend_send)

    logging_obj = MagicMock()
    logging_obj.litellm_trace_id = "test-trace-id"
    logging_obj.pre_call = MagicMock()
    logging_obj.async_success_handler = AsyncMock()
    logging_obj.success_handler = MagicMock()

    streaming = RealTimeStreaming(
        websocket=client_ws,
        backend_ws=backend_ws,
        logging_obj=logging_obj,
        provider_config=cfg,
        model="gemini-2.0-flash-live-001",
    )

    # Run backend→client forwarding for the three queued messages, then stop.
    # We don't run client_ack_messages here to avoid the blocking receive loop.
    await streaming.backend_to_client_send_messages()

    # --- Assertions ---

    # session.created should have been forwarded to client
    session_created_msgs = [m for m in sent_to_client if '"session.created"' in m]
    assert session_created_msgs, "Expected session.created to be sent to client"

    # At least one text delta should have been forwarded
    text_delta_msgs = [m for m in sent_to_client if '"response.output_text.delta"' in m]
    assert text_delta_msgs, "Expected response.output_text.delta to be sent to client"

    # Verify the delta contains the model's text
    delta_obj = json.loads(text_delta_msgs[0])
    assert "Hello from Vertex AI!" in delta_obj.get("delta", "")

    # response.done should have been forwarded
    done_msgs = [m for m in sent_to_client if '"response.done"' in m]
    assert done_msgs, "Expected response.done to be sent to client"


def test_vertex_warns_when_dropping_guardrail_turn_detection_update(caplog):
    """A subsequent session.update carrying the guardrail's
    ``create_response: False`` cannot be forwarded as a follow-up setup on
    Vertex AI (1007). Surface a warning so operators know the auto-response
    suppression is being silently dropped."""
    import logging

    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    session_update = {
        "type": "session.update",
        "session": {"turn_detection": {"create_response": False}},
    }

    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
        result = cfg.transform_realtime_request(
            json.dumps(session_update),
            "gemini-live-2.5-flash-preview-native-audio-09-2025",
            session_configuration_request=json.dumps({"setup": {"model": "x"}}),
        )

    assert result == []
    assert any(
        "Vertex AI Realtime" in record.message
        and "create_response=False" in record.message
        for record in caplog.records
    )


def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog):
    """A subsequent session.update without ``create_response: False`` is a
    routine drop and should stay at debug level (no warning)."""
    import logging

    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    session_update = {
        "type": "session.update",
        "session": {"instructions": "Be concise."},
    }

    with caplog.at_level(logging.WARNING, logger="LiteLLM"):
        cfg.transform_realtime_request(
            json.dumps(session_update),
            "gemini-live-2.5-flash-preview-native-audio-09-2025",
            session_configuration_request=json.dumps({"setup": {"model": "x"}}),
        )

    assert not any(
        "Vertex AI Realtime" in record.message and "session.update" in record.message
        for record in caplog.records
    )


@pytest.mark.asyncio
async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend(
    monkeypatch,
):
    """Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors.

    Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params``
    (the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here.
    """
    import websockets

    from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler

    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )

    captured: dict = {}

    def fake_connect(url, *args, **kwargs):
        captured["url"] = url
        raise RuntimeError("stop before establishing the backend connection")

    monkeypatch.setattr(websockets, "connect", fake_connect)

    try:
        await BaseLLMHTTPHandler().async_realtime(
            model="gemini-live-2.5-flash-preview-native-audio-09-2025",
            websocket=AsyncMock(),
            logging_obj=MagicMock(),
            provider_config=cfg,
            headers={},
            query_params={
                "model": "gemini-live-2.5-flash-preview-native-audio-09-2025",
                "intent": "chat",
            },
        )
    except (RuntimeError, Exception):
        pass

    assert "url" in captured, "websockets.connect was never called"
    assert "?" not in captured["url"]
    assert "model=" not in captured["url"]
    assert "intent=" not in captured["url"]


def test_vertex_function_call_output_omits_id():
    """Regression: Vertex Live rejects ``id`` on toolResponse.functionResponses (1007)."""
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    cfg._tool_call_id_to_name["call_abc123"] = "terminate_call"

    messages = cfg.transform_realtime_request(
        json.dumps(
            {
                "type": "conversation.item.create",
                "item": {
                    "type": "function_call_output",
                    "call_id": "call_abc123",
                    "output": '{"status": "ok"}',
                },
            }
        ),
        "gemini-live-2.5-flash-preview-native-audio-09-2025",
        session_configuration_request="existing",
    )

    assert len(messages) == 1
    payload = json.loads(messages[0])
    function_response = payload["toolResponse"]["functionResponses"][0]
    assert "id" not in function_response
    assert function_response["name"] == "terminate_call"
    assert function_response["response"] == {"status": "ok"}


def test_vertex_native_audio_keeps_requested_voice(patch_native_audio_cost_map_entry):
    """Regression: Vertex Live accepts speechConfig on native audio, so the client's voice must survive.

    Stripping it silently dropped voice selection for every Vertex native-audio
    session. TEXT is still coerced away, which Vertex does reject.
    """
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    session_update = {
        "type": "session.update",
        "session": {
            "output_modalities": ["text"],
            "audio": {"output": {"voice": "Aoede"}},
        },
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )

    generation_config = json.loads(messages[0])["setup"]["generationConfig"]
    assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Aoede"
    assert generation_config["responseModalities"] == ["AUDIO"]


def test_google_ai_studio_native_audio_keeps_requested_voice(patch_native_audio_cost_map_entry):
    """Regression: AI Studio native-audio Live accepts speechConfig too, so the voice survives on both providers."""
    from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig

    messages = GeminiRealtimeConfig().transform_realtime_request(
        json.dumps(
            {
                "type": "session.update",
                "session": {
                    "output_modalities": ["audio"],
                    "audio": {"output": {"voice": "Aoede"}},
                },
            }
        ),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )

    generation_config = json.loads(messages[0])["setup"]["generationConfig"]
    assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Aoede"


def test_vertex_native_audio_drops_openai_stock_voice(patch_native_audio_cost_map_entry):
    """Regression: OpenAI stock voice names must be dropped, not forwarded verbatim.

    Vertex Live closes the socket with 1007 on an unknown voice name, so a
    client sending OpenAI's default voice would lose the session entirely.
    Dropping the voice keeps the session alive on the model's default voice.
    """
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    session_update = {
        "type": "session.update",
        "session": {"audio": {"output": {"voice": "alloy"}}},
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )

    generation_config = json.loads(messages[0])["setup"]["generationConfig"]
    assert "speechConfig" not in generation_config


def test_vertex_native_audio_unmapped_voice_passes_through(patch_native_audio_cost_map_entry):
    """A voice name outside the OpenAI stock set is forwarded verbatim so Gemini-native names keep working."""
    cfg = VertexAIRealtimeConfig(
        access_token="tok", project="my-proj", location="us-central1"
    )
    session_update = {
        "type": "session.update",
        "session": {"audio": {"output": {"voice": "Kore"}}},
    }

    messages = cfg.transform_realtime_request(
        json.dumps(session_update),
        _NATIVE_AUDIO_MODEL,
        session_configuration_request=None,
    )

    generation_config = json.loads(messages[0])["setup"]["generationConfig"]
    assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
