"""
Integration tests for RealTimeStreaming guardrails against a live OpenAI backend.

These tests require OPENAI_API_KEY and are skipped if not set.

They verify end-to-end that:
  1. A text message blocked by a guardrail -> error event sent to client, the blocked
     message never reaches OpenAI, and the client's response.create is not forwarded.
  2. A voice transcript blocked by a guardrail -> error event sent, response.create NOT sent.
  3. A clean text message passes through and triggers a real OpenAI response.

Run with:
    uv run pytest tests/llm_translation/realtime/test_realtime_guardrails_openai.py -v -s
"""

import asyncio
import json
import os
from typing import List
from unittest.mock import AsyncMock, MagicMock

import pytest

import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
from litellm.types.guardrails import GuardrailEventHooks

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
OPENAI_REALTIME_URL = "wss://api.openai.com/v1/realtime?model=gpt-realtime"

pytestmark = pytest.mark.skipif(
    not OPENAI_API_KEY,
    reason="OPENAI_API_KEY not set - skipping OpenAI realtime integration tests",
)

# A unique phrase guaranteed NOT to appear in normal assistant output.
BLOCKED_PHRASE = "XSECRETBLOCKTESTPHRASEX"


class PhraseBlockingGuardrail(CustomGuardrail):
    """Blocks any message containing BLOCKED_PHRASE."""

    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
        for text in inputs.get("texts", []):
            if BLOCKED_PHRASE in text:
                raise ValueError("Content blocked: contains forbidden test phrase.")
        return inputs


def _make_guardrail(event_hook=GuardrailEventHooks.pre_call):
    return PhraseBlockingGuardrail(
        guardrail_name="integration-test-guard",
        event_hook=event_hook,
        default_on=True,
    )


class RecordingBackendWebSocket:
    """Wraps a real backend WebSocket and records every frame sent to it."""

    def __init__(self, backend_ws):
        self._backend_ws = backend_ws
        self.sent_messages: List[str] = []

    async def send(self, message):
        self.sent_messages.append(message)
        await self._backend_ws.send(message)

    async def recv(self, *args, **kwargs):
        return await self._backend_ws.recv(*args, **kwargs)

    async def close(self):
        await self._backend_ws.close()


async def _wait_for_event(client_events: List[dict], event_type: str, timeout: float = 15.0) -> dict:
    """Poll client_events list until an event with matching type appears."""
    deadline = asyncio.get_event_loop().time() + timeout
    while asyncio.get_event_loop().time() < deadline:
        matching = [e for e in client_events if e.get("type") == event_type]
        if matching:
            return matching[0]
        await asyncio.sleep(0.05)
    raise TimeoutError(f"Timed out waiting for '{event_type}'. Got so far: {[e.get('type') for e in client_events]}")


async def _build_streaming(client_events: List[dict], backend_ws, request_data=None):
    """Create a RealTimeStreaming with a mock client WebSocket that captures events."""
    client_ws = MagicMock()
    input_queue: asyncio.Queue = asyncio.Queue()

    async def send_text(data: str):
        client_events.append(json.loads(data))

    client_ws.send_text = send_text
    client_ws.receive_text = input_queue.get

    logging_obj = MagicMock()
    logging_obj.pre_call = MagicMock()
    logging_obj.async_success_handler = AsyncMock()
    logging_obj.success_handler = MagicMock()
    logging_obj.model_call_details = {}

    streaming = RealTimeStreaming(
        websocket=client_ws,
        backend_ws=backend_ws,
        logging_obj=logging_obj,
        request_data=request_data or {"guardrails": ["integration-test-guard"]},
    )
    return streaming, input_queue


@pytest.mark.asyncio
async def test_text_message_blocked_by_guardrail_no_ai_response():
    """
    Send a text message containing the blocked phrase, immediately followed by
    response.create (the reflexive client pattern).
    Guardrail must:
      - Send error event (guardrail_violation) to client.
      - Send response.output_audio_transcript.delta (or beta-protocol
        response.audio_transcript.delta) to client.
      - NEVER forward the blocked message to OpenAI.
      - Drop the client's response.create; the only response.create OpenAI sees
        is the guardrail's own (which voices the block message), so the model
        can never answer the blocked content.

    Assertions are on the recorded backend wire traffic, not on the model's
    reply wording: gpt-realtime phrases its voicing/refusal of the guardrail
    prompt nondeterministically, which made wording-based assertions flaky
    (see PRs #28191, #28200, #29477).
    """
    import websockets

    guardrail = _make_guardrail(GuardrailEventHooks.pre_call)
    litellm.callbacks = [guardrail]

    client_events: List[dict] = []

    try:
        async with websockets.connect(
            OPENAI_REALTIME_URL,
            additional_headers={
                "Authorization": f"Bearer {OPENAI_API_KEY}",
            },
        ) as raw_backend_ws:
            backend_ws = RecordingBackendWebSocket(raw_backend_ws)
            streaming, input_queue = await _build_streaming(client_events, backend_ws)

            backend_task = asyncio.create_task(streaming.backend_to_client_send_messages())
            client_task = asyncio.create_task(streaming.client_ack_messages())

            try:
                await _wait_for_event(client_events, "session.created", timeout=15)

                blocked_item = json.dumps(
                    {
                        "type": "conversation.item.create",
                        "item": {
                            "role": "user",
                            "content": [
                                {
                                    "type": "input_text",
                                    "text": f"Hello {BLOCKED_PHRASE}",
                                }
                            ],
                        },
                    }
                )
                await input_queue.put(blocked_item)
                await input_queue.put(json.dumps({"type": "response.create"}))

                await _wait_for_event(client_events, "response.done", timeout=30)

            finally:
                backend_task.cancel()
                client_task.cancel()
                await asyncio.gather(backend_task, client_task, return_exceptions=True)

        event_types = [e.get("type") for e in client_events]

        error_events = [e for e in client_events if e.get("type") == "error"]
        guardrail_errors = [e for e in error_events if e.get("error", {}).get("type") == "guardrail_violation"]
        assert len(guardrail_errors) >= 1, (
            f"Expected at least one guardrail_violation error but got: {[e.get('error', {}).get('type') for e in error_events]}"
        )

        transcript_deltas = [
            e
            for e in client_events
            if e.get("type")
            in (
                "response.output_audio_transcript.delta",
                "response.audio_transcript.delta",
            )
        ]
        assert len(transcript_deltas) >= 1, f"Expected guardrail message in transcript delta, got: {event_types}"

        sent_frames = backend_ws.sent_messages
        assert all(BLOCKED_PHRASE not in frame for frame in sent_frames), (
            f"Blocked message was forwarded to OpenAI: {sent_frames}"
        )

        sent_types = [json.loads(frame).get("type") for frame in sent_frames]
        assert sent_types.count("response.create") == 1, (
            f"Expected only the guardrail's response.create to reach OpenAI, got backend frames: {sent_types}"
        )
        assert sent_types.count("conversation.item.create") == 1, (
            f"Expected only the guardrail's conversation.item.create to reach OpenAI, got backend frames: {sent_types}"
        )

        done_events = [e for e in client_events if e.get("type") == "response.done"]
        assert len(done_events) >= 1, f"Expected response.done, got: {event_types}"
        for done in done_events:
            output = done.get("response", {}).get("output", [])
            ai_texts = [
                c.get("text", "") or c.get("transcript", "") for item in output for c in item.get("content", [])
            ]
            real_ai_text = " ".join(ai_texts).strip()
            assert BLOCKED_PHRASE not in real_ai_text, f"Blocked phrase leaked into AI response: {real_ai_text!r}"

    finally:
        litellm.callbacks = []


@pytest.mark.asyncio
async def test_voice_transcript_blocked_by_guardrail():
    """
    Simulate a backend-side voice transcription event containing the blocked phrase.
    Guardrail must block it - no response.create sent to OpenAI.
    """
    from websockets.exceptions import ConnectionClosed

    guardrail = _make_guardrail(GuardrailEventHooks.realtime_input_transcription)
    litellm.callbacks = [guardrail]

    client_events: List[dict] = []

    # Build the transcript event that would come from the OpenAI backend
    transcript_event = json.dumps(
        {
            "type": "conversation.item.input_audio_transcription.completed",
            "transcript": f"This is {BLOCKED_PHRASE} in my voice message",
            "item_id": "item_integ_test",
        }
    ).encode()

    # Mock backend that delivers the transcript then closes
    backend_ws = MagicMock()
    backend_ws.recv = AsyncMock(
        side_effect=[
            transcript_event,
            ConnectionClosed(None, None),
        ]
    )
    backend_ws.send = AsyncMock()

    try:
        streaming, _ = await _build_streaming(client_events, backend_ws)
        await streaming.backend_to_client_send_messages()

        event_types = [e.get("type") for e in client_events]

        # 1. Error event must be sent to client
        error_events = [e for e in client_events if e.get("type") == "error"]
        assert len(error_events) >= 1, f"Expected guardrail error event, got: {event_types}"
        assert error_events[0]["error"]["type"] == "guardrail_violation"

        # 2. Check what was sent to backend.
        #    The guardrail may send response.cancel + conversation.item.create (block msg)
        #    + response.create (to speak the block message). That's acceptable.
        #    What we assert is that a response.cancel was sent (blocking the original).
        sent_to_backend = [
            json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args and isinstance(c.args[0], str)
        ]
        response_cancels = [e for e in sent_to_backend if e.get("type") == "response.cancel"]
        assert len(response_cancels) >= 1 or len(sent_to_backend) == 0, (
            f"Guardrail should have sent response.cancel or nothing, got: {sent_to_backend}"
        )

        # Note: The guardrail may or may not send transcript deltas; the error event
        # (assertion #1) is the primary signal that the blocked content was handled.

    finally:
        litellm.callbacks = []


@pytest.mark.asyncio
async def test_clean_text_message_passes_through_to_openai():
    """
    A clean message (no blocked phrase) must pass the guardrail and result in a real
    AI response from OpenAI (response.done with non-empty output).
    """
    import websockets

    guardrail = _make_guardrail(GuardrailEventHooks.pre_call)
    litellm.callbacks = [guardrail]

    client_events: List[dict] = []

    try:
        async with websockets.connect(
            OPENAI_REALTIME_URL,
            additional_headers={
                "Authorization": f"Bearer {OPENAI_API_KEY}",
            },
        ) as backend_ws:
            streaming, input_queue = await _build_streaming(client_events, backend_ws)

            backend_task = asyncio.create_task(streaming.backend_to_client_send_messages())
            client_task = asyncio.create_task(streaming.client_ack_messages())

            try:
                await _wait_for_event(client_events, "session.created", timeout=15)

                # Send a clean message
                clean_item = json.dumps(
                    {
                        "type": "conversation.item.create",
                        "item": {
                            "role": "user",
                            "content": [{"type": "input_text", "text": "Reply with just: OK"}],
                        },
                    }
                )
                await input_queue.put(clean_item)
                await asyncio.sleep(0.1)
                await input_queue.put(json.dumps({"type": "response.create"}))

                # Wait for OpenAI to respond
                await _wait_for_event(client_events, "response.done", timeout=30)

            finally:
                backend_task.cancel()
                client_task.cancel()
                await asyncio.gather(backend_task, client_task, return_exceptions=True)

        # No guardrail error should have been sent
        error_events = [e for e in client_events if e.get("type") == "error"]
        guardrail_errors = [e for e in error_events if e.get("error", {}).get("type") == "guardrail_violation"]
        assert len(guardrail_errors) == 0, f"Clean message should not trigger guardrail, got: {guardrail_errors}"

        # AI response must be present
        done_events = [e for e in client_events if e.get("type") == "response.done"]
        assert len(done_events) >= 1, (
            f"Expected response.done from OpenAI, got: {[e.get('type') for e in client_events]}"
        )

    finally:
        litellm.callbacks = []
