"""EXAMPLE of a customer's OWN AUDIO SERVER for a Moholo embed (option 2, 2026-09-18): the visitor's browser connects here directly,
Moholo never carries the audio. Replace `reply()` with your speech-to-speech model; everything else is the protocol.

  pip install aiohttp numpy
  python3 own_audio_server_example.py --port 8797 [--wav reply.wav]     (no --wav: a built-in 2 s test sound is the reply, so it runs as downloaded)
  embed:   <script src="https://moholo.ai/embed.js" data-key="mk_..." data-audio="wss://voice.your.site/moholo"></script>
  or in the panel: "my own audio server" + the wss:// address.

PROTOCOL (WebSocket, one connection per visitor)
  browser -> you   JSON {type:'hello', avatar, key, visitor, lang, in:{format:'s16',rate:16000,channels:1}, out:{format:'s16',rate:24000,channels:1}}
                   binary  one whole spoken utterance: int16 mono 16 kHz PCM   (followed by JSON {type:'utterance', rate, seconds})
                   JSON {type:'text', text}               the visitor typed
  you -> browser   binary  reply audio chunks: int16 mono 24 kHz PCM (any chunk size; 20-200 ms is good) - lips move as they arrive
                   JSON {type:'audio', pcm:<base64 int16>, rate?, end?}   alternative to binary
                   JSON {type:'end'}                      the reply is complete (required, ends the turn)
                   JSON {type:'stop'}                     interrupt what is playing
                   JSON {type:'caption', text}            optional caption under the avatar
                   JSON {type:'status', text}             optional status line
                   JSON {type:'format', rate}             optional: change the reply sample rate (default 24000)
"""
import argparse, asyncio, json, wave, time
import numpy as np
from aiohttp import web, WSMsgType

def synth_reply(rate_out=24000, seconds=2.0):
    """No --wav given: a speech-like test sound (a 140 Hz voice-ish tone under a 4 Hz syllable envelope), so the mouth visibly moves."""
    t = np.arange(int(rate_out * seconds)) / rate_out; env = 0.5 - 0.5 * np.cos(2 * np.pi * 4 * t)
    x = env * (0.6 * np.sin(2 * np.pi * 140 * t) + 0.3 * np.sin(2 * np.pi * 280 * t) + 0.15 * np.sin(2 * np.pi * 700 * t))
    return (x * 12000).astype(np.int16)

def load_reply(path, rate_out=24000):
    if not path: return synth_reply(rate_out)
    import os
    if not os.path.isfile(path): raise SystemExit(f'--wav {path}: no such file (give a 16-bit PCM .wav, or leave --wav out for the built-in test sound)')
    w = wave.open(path); sr = w.getframerate(); n = w.getnframes(); x = np.frombuffer(w.readframes(n), dtype=np.int16).astype(np.float32)
    if w.getnchannels() == 2: x = x.reshape(-1, 2).mean(1)
    if sr != rate_out:   # linear resample (good enough for an example)
        t_in = np.arange(len(x)) / sr; t_out = np.arange(int(len(x) * rate_out / sr)) / rate_out; x = np.interp(t_out, t_in, x)
    return x.astype(np.int16)

async def reply(ws, prompt_kind, reply_pcm, rate=24000, chunk_ms=100):
    """Stream the reply audio in chunks (a real server streams its model output as it is generated)."""
    await ws.send_str(json.dumps({'type': 'caption', 'text': f'(example reply to your {prompt_kind})'}))
    step = rate * chunk_ms // 1000; t0 = time.time()
    for i in range(0, len(reply_pcm), step):
        await ws.send_bytes(reply_pcm[i:i + step].tobytes())
        await asyncio.sleep(max(0, (i + step) / rate - (time.time() - t0)))   # real time, like a streaming TTS
    await ws.send_str(json.dumps({'type': 'end'}))

async def handler(request):
    ws = web.WebSocketResponse(max_msg_size=32 * 1024 * 1024); await ws.prepare(request); reply_pcm = request.app['reply']
    print('[own-audio] visitor connected')
    async for m in ws:
        if m.type == WSMsgType.BINARY:
            pcm = np.frombuffer(m.data, dtype=np.int16); print(f'[own-audio] utterance {len(pcm) / 16000:.2f} s -> replying')
            asyncio.create_task(reply(ws, 'speech', reply_pcm))
        elif m.type == WSMsgType.TEXT:
            try: d = json.loads(m.data)
            except Exception: continue
            if d.get('type') == 'hello': print('[own-audio] hello', {k: d.get(k) for k in ('avatar', 'key', 'lang')})
            elif d.get('type') == 'text': print('[own-audio] text:', d.get('text')); asyncio.create_task(reply(ws, 'text', reply_pcm))
    print('[own-audio] visitor left'); return ws

if __name__ == '__main__':
    ap = argparse.ArgumentParser(); ap.add_argument('--port', type=int, default=8797); ap.add_argument('--wav', default='', help='16-bit PCM wav played as every reply; omitted = built-in test sound'); a = ap.parse_args()
    app = web.Application(); app['reply'] = load_reply(a.wav); app.router.add_get('/ws', handler)
    print(f'[own-audio] example server ws://0.0.0.0:{a.port}/ws  reply={len(app["reply"]) / 24000:.1f} s'); web.run_app(app, port=a.port, print=None)