"""
Simple xAI Voice Agent using LiveKit SDK with LiteLLM Gateway

This example shows how to use LiveKit's xAI realtime plugin through LiteLLM proxy.
LiteLLM acts as a unified interface, allowing you to switch between xAI, OpenAI, 
and Azure realtime APIs without changing your agent code.
"""

import asyncio
import json
import os
import websockets

# Configuration
PROXY_URL = os.getenv("LITELLM_PROXY_URL", "http://localhost:4000")
API_KEY = os.getenv("LITELLM_API_KEY", "sk-1234")
MODEL = os.getenv("LITELLM_MODEL", "grok-voice-agent")


async def run_voice_agent():
    """
    Simple voice agent that:
    1. Connects to xAI realtime API through LiteLLM proxy
    2. Sends a user message
    3. Streams back the response
    """

    url = f"ws://{PROXY_URL.replace('http://', '').replace('https://', '')}/v1/realtime?model={MODEL}"
    headers = {"Authorization": f"Bearer {API_KEY}"}

    print(f"🎙️  Connecting to voice agent...")
    print(f"   Model: {MODEL}")
    print(f"   Proxy: {PROXY_URL}")
    print()

    async with websockets.connect(url, additional_headers=headers) as ws:
        # Receive initial connection event
        initial = json.loads(await ws.recv())
        print(f"✅ Connected! Event: {initial['type']}\n")

        # Get user input
        user_message = input("💬 Your message: ").strip()
        if not user_message:
            user_message = "Tell me a fun fact about AI!"

        print(f"\n🤖 Sending to {MODEL}...\n")

        # Send user message
        await ws.send(
            json.dumps(
                {
                    "type": "conversation.item.create",
                    "item": {
                        "type": "message",
                        "role": "user",
                        "content": [{"type": "input_text", "text": user_message}],
                    },
                }
            )
        )

        # Request response
        await ws.send(
            json.dumps(
                {
                    "type": "response.create",
                    "response": {"modalities": ["text", "audio"]},
                }
            )
        )

        # Stream response
        print("🎤 Response: ", end="", flush=True)
        transcript = []

        try:
            while True:
                msg = await asyncio.wait_for(ws.recv(), timeout=15.0)
                event = json.loads(msg)

                # Capture transcript deltas
                if event["type"] == "response.output_audio_transcript.delta":
                    delta = event.get("delta", "")
                    if delta:
                        print(delta, end="", flush=True)
                        transcript.append(delta)

                # Done when response completes
                elif event["type"] == "response.done":
                    break

        except asyncio.TimeoutError:
            pass

        print("\n")

        if transcript:
            print(f"✅ Complete response: {''.join(transcript)}")

        await ws.close()


def main():
    """Run the voice agent"""
    print("=" * 70)
    print("LiveKit xAI Voice Agent via LiteLLM Proxy")
    print("=" * 70)
    print()

    try:
        asyncio.run(run_voice_agent())
    except KeyboardInterrupt:
        print("\n\n👋 Goodbye!")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        print("\nMake sure LiteLLM proxy is running:")
        print(f"  litellm --config config.yaml --port 4000")


if __name__ == "__main__":
    main()