#!/usr/bin/env python3
"""OpenAI Realtime reference brain -- Python implementation (slice-3 spec §7.5).

Mirrors rutster-brain-realtime's WS-server + OpenAI-WS-client glue in
Python, proving the slice-3 tap protocol is language-agnostic. The mapping
is the §4.2 event table verbatim.

Run:
    pip install -r requirements.txt
    OPENAI_API_KEY=sk-... python examples/openai_realtime_brain/openai_realtime_brain.py

Not in CI (violates the zero-non-Rust-dev-deps dev loop per AGENTS.md).
"""

import asyncio
import json
import os
import sys

import websockets

PROTOCOL_VERSION = 1
SAMPLES = 480

# Configuration via env (matching the Rust binary's naming convention so the
# docs read the same on both sides).
RUTSTER_TAP_BIND = os.environ.get("RUTSTER_BRAIN_BIND", "127.0.0.1:8082")
OPENAI_MODEL = os.environ.get("OPENAI_REALTIME_MODEL", "gpt-4o-realtime")
OPENAI_VOICE = os.environ.get("OPENAI_REALTIME_VOICE", "alloy")


def openai_realtime_url(model: str) -> str:
    """wss://api.openai.com/v1/realtime?model=<model> (spec §4.2 + §5.3)."""
    return f"wss://api.openai.com/v1/realtime?model={model}"


def openai_headers(api_key: str):
    """The two headers the OpenAI WS handshake requires (spec §5.3)."""
    return [
        ("Authorization", f"Bearer {api_key}"),
        ("OpenAI-Beta", "realtime=v1"),
    ]


async def handle_tap_connection(tap_ws):
    """Bridge one rutster tap WS connection to one OpenAI Realtime session.

    Spec §4.2 mapping table, in both directions. One tap connection == one
    OpenAI Realtime session; stateless across reconnects (spec §5.3), exactly
    like the Rust brain and the slice-2 echo brain pattern.
    """
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        await tap_ws.close(code=1011, reason="OPENAI_API_KEY required")
        return

    # Open the OpenAI WS with the two required headers. additional_headers is
    # the websockets>=12 spelling (was extra_headers in v11).
    openai_ws = await websockets.connect(
        openai_realtime_url(OPENAI_MODEL),
        additional_headers=openai_headers(api_key),
    )

    # S4: send session.update with turn_detection: null. The FOB (the core)
    # owns turn-taking; OpenAI's server-side VAD is explicitly disabled so it
    # does not contradict core-authoritative playout decisions. See
    # openai_client.rs for the Rust-side commentary on why this matters.
    await openai_ws.send(json.dumps({
        "type": "session.update",
        "session": {
            "modalities": ["text", "audio"],
            "voice": OPENAI_VOICE,
            "input_audio_format": "pcm16",
            "output_audio_format": "pcm16",
            "sample_rate": 24000,
            "turn_detection": None,
        },
    }))
    print(f"[openai_brain] session.update sent (turn_detection=null); "
          f"model={OPENAI_MODEL} voice={OPENAI_VOICE}")

    seq_egress = 0

    def next_seq():
        nonlocal seq_egress
        s = seq_egress
        seq_egress += 1
        return s

    async def tap_to_openai():
        """rutster tap events -> OpenAI Realtime events (spec §4.2 left column)."""
        async for raw in tap_ws:
            try:
                env = json.loads(raw)
            except json.JSONDecodeError as e:
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "error", "seq": next_seq(),
                    "ts": 0, "code": "decode_failed", "message": str(e),
                }))
                continue

            t = env.get("type")
            if t == "hello":
                # Stateless across reconnects: ack hello, do NOT re-send
                # session.update (already sent on this OpenAI WS session).
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "hello", "seq": next_seq(),
                    "ts": 0, "session_id": env.get("session_id", ""),
                }))
            elif t == "audio_in":
                # Pass-through: same base64 PCM wire shape. No transcoding.
                await openai_ws.send(json.dumps({
                    "type": "input_audio_buffer.append",
                    "audio": env["pcm"],
                }))
            elif t == "function_call_output":
                # Core's tool-registry reply -> OpenAI conversation.item.create
                # so the model can continue. Spec §4.2 row 5.
                await openai_ws.send(json.dumps({
                    "type": "conversation.item.create",
                    "item": {
                        "type": "function_call_output",
                        "call_id": env["id"],
                        "output": json.dumps(env.get("result", {})),
                    },
                }))
            elif t in ("bye", "session_end"):
                await openai_ws.close()
                return
            # Unknown frame types: ignore (forward-compat via additive protocol).

    async def openai_to_tap():
        """OpenAI Realtime events -> rutster tap events (spec §4.2 right column)."""
        async for raw in openai_ws:
            try:
                event = json.loads(raw)
            except json.JSONDecodeError:
                continue
            t = event.get("type")
            if t == "response.audio.delta":
                # Forward as tap audio_out. The delta is base64 PCM16; pass
                # through. (Batching into 20 ms frames is a future
                # optimization, spec §1.2; v1 passes each delta through.)
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "audio_out",
                    "seq": next_seq(), "ts": 0,
                    "pcm": event["delta"], "samples": SAMPLES,
                }))
            elif t == "input_audio_buffer.speech_started":
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "speech_started",
                    "seq": next_seq(), "ts": 0,
                }))
            elif t == "input_audio_buffer.speech_stopped":
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "speech_stopped",
                    "seq": next_seq(), "ts": 0,
                }))
            elif t == "response.function_call_arguments.done":
                # Tool invocation from the model -> core's tool registry.
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "function_call",
                    "id": event["call_id"], "name": event["name"],
                    "args": json.loads(event.get("arguments", "{}")),
                    "seq": next_seq(), "ts": 0,
                }))
            elif t == "error":
                # Surface OpenAI-side errors to the core for observation.
                await tap_ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "error",
                    "seq": next_seq(), "ts": 0,
                    "code": "openai_error",
                    "message": event.get("error", {}).get("message", "unknown"),
                }))

    try:
        await asyncio.gather(tap_to_openai(), openai_to_tap())
    finally:
        await openai_ws.close()


async def main():
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        sys.exit("OPENAI_API_KEY required (see README.md)")

    host, port = RUTSTER_TAP_BIND.split(":")
    print(f"[openai_brain] listening on ws://{RUTSTER_TAP_BIND} "
          f"(OpenAI model={OPENAI_MODEL} voice={OPENAI_VOICE})")
    async with websockets.serve(handle_tap_connection, host, int(port)):
        await asyncio.Future()  # run forever


if __name__ == "__main__":
    asyncio.run(main())
