#!/usr/bin/env python3
"""Rutster slice-2 reference echo brain (foreign-language canonical demo).

Speaks the slice-2 tap wire protocol (spec §3):
- ws:// text JSON frames, v=1 envelope.
- On `hello`: ack with `hello`.
- On `audio_in`: echo the same PCM back as `audio_out`.
- On `bye` / `session_end`: close cleanly.

Run:
    pip install websockets
    python examples/echo_brain/echo_brain.py
"""

import asyncio
import json
import base64

import websockets

PROTOCOL_VERSION = 1
SAMPLES = 480


async def echo_handler(ws):
    seq_egress = 0
    # Wait for hello.
    raw = await ws.recv()
    env = json.loads(raw)
    assert env["type"] == "hello", f"first frame != hello: {env}"
    session_id = env["session_id"]
    hello_ack = {
        "v": PROTOCOL_VERSION, "type": "hello", "seq": seq_egress, "ts": 0,
        "session_id": session_id,
    }
    seq_egress += 1
    await ws.send(json.dumps(hello_ack))
    print(f"[echo_brain] hello acked: session_id={session_id}")

    async for raw in ws:
        try:
            env = json.loads(raw)
        except json.JSONDecodeError as e:
            await ws.send(json.dumps({
                "v": PROTOCOL_VERSION, "type": "error", "seq": seq_egress, "ts": 0,
                "code": "decode_failed", "message": str(e),
            }))
            seq_egress += 1
            continue

        kind = env.get("type")
        if kind == "audio_in":
            pcm_b64 = env["pcm"]
            samples = env.get("samples", SAMPLES)
            if samples != SAMPLES:
                await ws.send(json.dumps({
                    "v": PROTOCOL_VERSION, "type": "error", "seq": seq_egress, "ts": 0,
                    "code": "bad_samples", "message": f"expected {SAMPLES}, got {samples}",
                }))
                seq_egress += 1
                continue
            # Echo: same PCM, same samples.
            out = {
                "v": PROTOCOL_VERSION, "type": "audio_out", "seq": seq_egress,
                "ts": env.get("ts", 0),
                "pcm": pcm_b64, "samples": samples,
            }
            seq_egress += 1
            await ws.send(json.dumps(out))
        elif kind == "bye":
            print(f"[echo_brain] bye received: {env.get('reason')}")
            await ws.send(json.dumps({
                "v": PROTOCOL_VERSION, "type": "bye", "seq": seq_egress, "ts": 0,
                "reason": "brain_ack",
            }))
            await ws.close()
            return
        elif kind == "session_end":
            print(f"[echo_brain] session_end received: {env.get('reason')}")
            await ws.close()
            return
        else:
            print(f"[echo_brain] unknown frame type: {kind}; ignoring")


async def main():
    print("[echo_brain] listening on ws://127.0.0.1:8081/echo")
    async with websockets.serve(echo_handler, "127.0.0.1", 8081):
        await asyncio.Future()  # run forever


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