#!/usr/bin/env python3
"""WebSocket proxy: browser <-> Google Gemini Live API.

The browser never sees the API key — it connects to ws://127.0.0.1:8191/
and this proxy opens its own upstream WebSocket to Google with the secret
key from `.env`. All messages are forwarded as-is in both directions.

Bound to 127.0.0.1 only. Run alongside serve.py (or it gets auto-spawned).
"""
import asyncio
import json
import os
import sys
from pathlib import Path

import websockets

# ── .env loader (tiny, no extra deps) ────────────────────────────────────
def load_dotenv():
    env_path = Path(__file__).parent / ".env"
    if not env_path.exists():
        return
    for raw in env_path.read_text().splitlines():
        s = raw.strip()
        if not s or s.startswith("#") or "=" not in s:
            continue
        k, v = s.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip())


load_dotenv()

API_KEY = os.environ.get("GEMINI_API_KEY", "").strip()
DEFAULT_MODEL = os.environ.get(
    "GEMINI_LIVE_MODEL", "gemini-2.5-flash-native-audio-latest"
).strip()
HOST = "127.0.0.1"
PORT = int(os.environ.get("WS_PORT", "3001"))

if not API_KEY:
    print(
        "[ws-proxy] FATAL: GEMINI_API_KEY missing. Put it in v3/.env",
        file=sys.stderr,
    )
    sys.exit(2)

GEMINI_URL = (
    "wss://generativelanguage.googleapis.com/ws/"
    "google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent"
    f"?key={API_KEY}"
)


async def proxy_connection(browser_ws):
    """One browser <-> one upstream Gemini Live session. Bidirectional copy."""
    peer = getattr(browser_ws, "remote_address", ("?",))[0]
    print(f"[ws-proxy] + client {peer}", flush=True)
    try:
        async with websockets.connect(
            GEMINI_URL,
            max_size=2**24,           # 16 MiB — accommodate big audio chunks
            ping_interval=20,
            ping_timeout=20,
        ) as upstream:

            async def browser_to_upstream():
                async for msg in browser_ws:
                    await upstream.send(msg)

            async def upstream_to_browser():
                async for msg in upstream:
                    await browser_ws.send(msg)

            done, pending = await asyncio.wait(
                [
                    asyncio.create_task(browser_to_upstream()),
                    asyncio.create_task(upstream_to_browser()),
                ],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
    except websockets.exceptions.ConnectionClosed:
        pass
    except Exception as e:
        print(f"[ws-proxy] {peer} err: {type(e).__name__}: {e}", flush=True)
        try:
            await browser_ws.close(code=1011, reason=str(e))
        except Exception:
            pass
    finally:
        print(f"[ws-proxy] - client {peer}", flush=True)


async def main():
    print(
        f"[ws-proxy] listening on ws://{HOST}:{PORT}  model={DEFAULT_MODEL}",
        flush=True,
    )
    async with websockets.serve(
        proxy_connection,
        HOST,
        PORT,
        max_size=2**24,
        ping_interval=20,
    ):
        await asyncio.Future()  # run forever


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        sys.exit(0)