tests/scripts/wireless_e2e.py

"""End-to-end test of the upload+play+cancel commands against the
live wireless daemon over WebSocket.

Mimics what a client app does, but from Python so we can iterate
quickly without needing a browser. Uses a raw websockets client so we
see the broadcast events (the SDK's typed WSClient filters everything
that doesn't match its known message types).

Wire shape (fire-and-forget):
  - upload_move_start/chunk/finish: no acks; just fire
  - play_uploaded_move: no per-call ack; daemon broadcasts two
    events tagged type="play_uploaded_move" and upload_id=<id>:
      * {started: true, duration_s: D}  when the loop starts
      * {finished: true} | {cancelled: true} | {error: "..."}
  - cancel_move: fire-and-forget

Test cases:
  1. happy path  -- upload intro.json + play to completion
  2. cancel mid-play  -- upload + play + cancel after 0.5 s
  3. chunk sizes  -- upload a long synthetic move (30 s @ 100 Hz)
                     in 12 KB chunks, confirm no failure

Reports timings: chunk send pace, time-to-started, time-to-final,
upload throughput.
"""

import asyncio
import base64
import gzip
import json
import math
import sys
import time
import uuid
from pathlib import Path

import requests
import websockets

DAEMON_HOST = "reachy-mini.local"
DAEMON_PORT = 8000
WS_URL = f"ws://{DAEMON_HOST}:{DAEMON_PORT}/ws/sdk"
HTTP_URL = f"http://{DAEMON_HOST}:{DAEMON_PORT}"

CHUNK_SIZE = 12 * 1024


async def wait_for(ws, predicate, timeout, debug=""):
    deadline = time.perf_counter() + timeout
    while True:
        remaining = deadline - time.perf_counter()
        if remaining <= 0:
            raise TimeoutError(f"wait_for timed out: {debug}")
        raw = await asyncio.wait_for(ws.recv(), timeout=remaining)
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if predicate(msg):
            return msg


async def upload_move(ws, move_dict, description="test", encoding="json"):
    """Fire start + all chunks + finish, no awaits in between.

    encoding="json"        -- raw UTF-8 JSON chunks (smallest code path)
    encoding="gzip+base64" -- gzip+base64 the full JSON, then chunk
    """
    raw_json = json.dumps(move_dict)
    if encoding == "gzip+base64":
        compressed = gzip.compress(raw_json.encode("utf-8"))
        payload = base64.b64encode(compressed).decode("ascii")
    elif encoding == "json":
        payload = raw_json
    else:
        raise ValueError(f"unknown encoding {encoding!r}")
    chunks = [payload[i:i + CHUNK_SIZE] for i in range(0, len(payload), CHUNK_SIZE)]
    upload_id = f"py-{uuid.uuid4().hex[:8]}"

    print(f"  uploading move ({encoding}): "
          f"raw_json={len(raw_json)} bytes, wire={len(payload)} bytes "
          f"({len(payload)/len(raw_json)*100:.0f}%), {len(chunks)} chunks "
          f"upload_id={upload_id}")

    t0 = time.perf_counter()
    await ws.send(json.dumps({
        "type": "upload_move_start",
        "upload_id": upload_id,
        "total_chunks": len(chunks),
        "description": description,
        "estimated_duration_s": 0.0,
        "encoding": encoding,
    }))
    for i, c in enumerate(chunks):
        await ws.send(json.dumps({
            "type": "upload_move_chunk",
            "upload_id": upload_id,
            "chunk_index": i,
            "chunk": c,
        }))
    await ws.send(json.dumps({
        "type": "upload_move_finish",
        "upload_id": upload_id,
    }))
    elapsed = time.perf_counter() - t0
    rate = len(payload) / elapsed / 1024 if elapsed > 0 else 0
    print(f"    {len(chunks) + 2} ws.send() calls in {elapsed*1000:.0f} ms "
          f"({rate:.1f} KB/s on wire)")
    return upload_id


async def upload_audio(ws, wav_bytes, upload_id):
    """Fire start + chunks + finish for an audio WAV. Fire-and-forget."""
    payload = base64.b64encode(wav_bytes).decode("ascii")
    chunks = [payload[i:i + CHUNK_SIZE] for i in range(0, len(payload), CHUNK_SIZE)]
    print(f"  audio: raw={len(wav_bytes)} bytes, base64={len(payload)} bytes, {len(chunks)} chunks")
    t0 = time.perf_counter()
    await ws.send(json.dumps({
        "type": "upload_audio_start",
        "upload_id": upload_id,
        "total_chunks": len(chunks),
        "encoding": "wav-base64",
        "description": "audio for " + upload_id,
    }))
    for i, c in enumerate(chunks):
        await ws.send(json.dumps({
            "type": "upload_audio_chunk",
            "upload_id": upload_id,
            "chunk_index": i,
            "chunk": c,
        }))
    await ws.send(json.dumps({
        "type": "upload_audio_finish",
        "upload_id": upload_id,
    }))
    elapsed = time.perf_counter() - t0
    print(f"    audio upload: {elapsed*1000:.0f} ms "
          f"({len(payload)/max(elapsed,1e-6)/1024:.1f} KB/s on wire)")


async def play(ws, upload_id, *, play_frequency=100.0, cancel_after_s=None,
               audio_lead_ms=0.0):
    """Send play_uploaded_move, await broadcast started + final events."""
    t_send = time.perf_counter()
    await ws.send(json.dumps({
        "type": "play_uploaded_move",
        "upload_id": upload_id,
        "play_frequency": play_frequency,
        "initial_goto_duration": 0.0,
        "audio_lead_ms": audio_lead_ms,
    }))
    started = await wait_for(
        ws,
        lambda m: (m.get("type") == "play_uploaded_move"
                   and m.get("upload_id") == upload_id
                   and m.get("started") is True),
        timeout=8.0,
        debug="play_uploaded_move started",
    )
    t_started = time.perf_counter() - t_send
    duration_s = started.get("duration_s", 0.0)
    has_audio = started.get("has_audio", False)
    print(f"    started event: {t_started*1000:.1f} ms after play sent "
          f"(declared duration {duration_s:.2f}s, has_audio={has_audio})")

    if cancel_after_s is not None:
        await asyncio.sleep(cancel_after_s)
        print(f"    sending cancel_move at t+{cancel_after_s:.2f}s")
        await ws.send(json.dumps({"type": "cancel_move"}))

    final = await wait_for(
        ws,
        lambda m: (m.get("type") == "play_uploaded_move"
                   and m.get("upload_id") == upload_id
                   and (m.get("finished") is True
                        or m.get("cancelled") is True
                        or "error" in m)),
        timeout=duration_s + 30,
        debug="play_uploaded_move final",
    )
    total = time.perf_counter() - t_send
    print(f"    final event: {total:.2f}s after play sent -- {final}")
    return final


def build_synthetic_move(duration_s, hz=100):
    """A simple sweep: head x oscillates +-0.05 m sinusoidally."""
    frames = []
    times = []
    n = int(duration_s * hz)
    for i in range(n):
        t = i / hz
        times.append(round(t, 4))
        x = 0.05 * math.sin(2 * math.pi * 0.5 * t)
        frames.append({
            "head": [
                [1, 0, 0, x],
                [0, 1, 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1],
            ],
            "antennas": [0.0, 0.0],
            "body_yaw": 0.0,
        })
    return {
        "description": f"synthetic sweep {duration_s}s @ {hz}Hz",
        "time": times,
        "set_target_data": frames,
    }


async def main():
    print("Reset to base...")
    requests.post(f"{HTTP_URL}/api/move/goto",
                  json={"head_pose": {"x": 0, "y": 0, "z": 0, "roll": 0, "pitch": 0, "yaw": 0},
                        "antennas": [-0.1745, 0.1745], "duration": 0.8},
                  timeout=5)
    await asyncio.sleep(1.5)

    print(f"Connecting to {WS_URL}")
    async with websockets.connect(WS_URL, max_size=2**24) as ws:
        # Test 1
        print("\n=== Test 1: intro.json happy path ===")
        with open("/Users/remi/reachy_mini_apps/marionette-experimental/intro.json") as f:
            intro = json.load(f)
        upload_id = await upload_move(ws, intro, description="intro")
        await play(ws, upload_id)
        await asyncio.sleep(1.5)
        requests.post(f"{HTTP_URL}/api/move/goto",
                      json={"head_pose": {"x": 0, "y": 0, "z": 0, "roll": 0, "pitch": 0, "yaw": 0},
                            "antennas": [-0.1745, 0.1745], "duration": 0.8},
                      timeout=5)
        await asyncio.sleep(1.5)

        # Test 2
        print("\n=== Test 2: cancel mid-play ===")
        upload_id = await upload_move(ws, intro, description="intro-cancel")
        await play(ws, upload_id, cancel_after_s=0.7)
        await asyncio.sleep(1.5)
        requests.post(f"{HTTP_URL}/api/move/goto",
                      json={"head_pose": {"x": 0, "y": 0, "z": 0, "roll": 0, "pitch": 0, "yaw": 0},
                            "antennas": [-0.1745, 0.1745], "duration": 0.8},
                      timeout=5)
        await asyncio.sleep(1.5)

        # Test 3
        print("\n=== Test 3: 30 s synthetic move, json encoding ===")
        long_move = build_synthetic_move(30.0, hz=100)
        upload_id = await upload_move(ws, long_move, description="30s-json",
                                      encoding="json")
        await play(ws, upload_id, cancel_after_s=1.0)
        await asyncio.sleep(1.0)

        # Test 4
        print("\n=== Test 4: 30 s synthetic move, gzip+base64 encoding ===")
        upload_id = await upload_move(ws, long_move, description="30s-gz64",
                                      encoding="gzip+base64")
        await play(ws, upload_id, cancel_after_s=1.0)
        await asyncio.sleep(1.0)

        # Test 5 (the target case: a 3 minute move)
        print("\n=== Test 5: 3 minute synthetic move, gzip+base64 ===")
        long_move = build_synthetic_move(180.0, hz=100)
        upload_id = await upload_move(ws, long_move, description="3min-gz64",
                                      encoding="gzip+base64")
        await play(ws, upload_id, cancel_after_s=1.0)
        await asyncio.sleep(1.0)

        # Test 6: intro.json + audio (single-clock daemon-side sync)
        print("\n=== Test 6: move + audio, daemon plays both ===")
        # Find any small WAV in the extract dir, or fall back to
        # generating ~3 s of 16 kHz mono PCM silence (just to exercise
        # the path).
        wav_path = "/Users/remi/Downloads/anne-charlotte-music-extract/the-white-stripes-seven-nation-army.wav"
        try:
            with open(wav_path, "rb") as f:
                wav_bytes = f.read()
            print(f"  using {wav_path}")
        except FileNotFoundError:
            print("  WAV file not found; synthesizing 3 s of silence as placeholder")
            wav_bytes = build_silence_wav(seconds=3.0, sample_rate=16000)
        upload_id = await upload_move(ws, intro, description="intro+audio",
                                      encoding="gzip+base64")
        await upload_audio(ws, wav_bytes, upload_id)
        await play(ws, upload_id, audio_lead_ms=0.0)
        await asyncio.sleep(1.0)
        requests.post(f"{HTTP_URL}/api/move/goto",
                      json={"head_pose": {"x": 0, "y": 0, "z": 0, "roll": 0, "pitch": 0, "yaw": 0},
                            "antennas": [-0.1745, 0.1745], "duration": 0.8},
                      timeout=5)
        await asyncio.sleep(1.5)

        # Test 7: cancel mid-play of a move with audio
        print("\n=== Test 7: cancel mid-play with audio ===")
        upload_id = await upload_move(ws, intro, description="intro+audio-cancel",
                                      encoding="gzip+base64")
        await upload_audio(ws, wav_bytes, upload_id)
        await play(ws, upload_id, cancel_after_s=1.0)

    print("\nAll tests done.")


def build_silence_wav(seconds: float, sample_rate: int = 16000) -> bytes:
    """Minimal mono 16-bit PCM WAV with all-zero samples."""
    import struct
    n_samples = int(seconds * sample_rate)
    data_size = n_samples * 2
    return (
        b"RIFF" + struct.pack("<I", 36 + data_size) + b"WAVE"
        + b"fmt " + struct.pack("<IHHIIHH", 16, 1, 1, sample_rate, sample_rate * 2, 2, 16)
        + b"data" + struct.pack("<I", data_size)
        + b"\x00" * data_size
    )


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