Update app.py

app.py

@@ -1,34 +1,116 @@
-"""Reachy Mini Controller - Fun, Queue-Based Control Interface"""
+"""
+Reachy Mini Controller - Cleaned Version
+A centralized server that listens for Robot connections and hosts a Gradio control interface.
+"""
 
 import asyncio
-import subprocess
-import sys
-import threading
+import io
 import json
+import threading
 import time
+import queue
 from dataclasses import dataclass
 from typing import List, Optional
 
 import cv2
 import gradio as gr
 import numpy as np
-from fastapi import FastAPI, WebSocket
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import StreamingResponse
+import uvicorn
+from fastrtc import WebRTC
+
+# Try to import the utility, handle error if running in standalone test without library
+try:
+    from reachy_mini.utils import create_head_pose
+except ImportError:
+    print("⚠️ Warning: reachy_mini module not found. Mocking create_head_pose for testing.")
+    def create_head_pose(**kwargs): return np.array([0,0,0])
+
+AUDIO_SAMPLE_RATE = 16000  # respeaker samplerate
+
+import struct
+
+def gen_wav_header(sample_rate, bits_per_sample, channels):
+    """
+    Generates a generic WAV header. 
+    We set the file size to 0xFFFFFFFF (max) to trick the browser into 
+    thinking it's a very long file for streaming purposes.
+    """
+    datasize = 0xFFFFFFFF
+    o = bytes("RIFF", 'ascii')  # (4byte) Marks file as RIFF
+    o += struct.pack('<I', datasize + 36)  # (4byte) File size in bytes excluding "RIFF" and size
+    o += bytes("WAVE", 'ascii')  # (4byte) File type
+    o += bytes("fmt ", 'ascii')  # (4byte) Format Chunk Marker
+    o += struct.pack('<I', 16)  # (4byte) Length of above format data
+    o += struct.pack('<H', 1)  # (2byte) Format type (1 - PCM)
+    o += struct.pack('<H', channels)  # (2byte)
+    o += struct.pack('<I', sample_rate)  # (4byte)
+    o += struct.pack('<I', sample_rate * channels * bits_per_sample // 8)  # (4byte)
+    o += struct.pack('<H', channels * bits_per_sample // 8)  # (2byte)
+    o += struct.pack('<H', bits_per_sample)  # (2byte)
+    o += bytes("data", 'ascii')  # (4byte) Data Chunk Marker
+    o += struct.pack('<I', datasize)  # (4byte) Data size in bytes
+    return o
+
+# --- 1. Global State Management ---
+class GlobalState:
+    """
+    Singleton-style class to manage shared state between FastAPI (WebSockets)
+    and Gradio (UI Thread).
+    """
+    def __init__(self):
+        # Connection handles
+        self.robot_ws: Optional[WebSocket] = None
+        self.robot_loop: Optional[asyncio.AbstractEventLoop] = None
+        
+        # Video Stream Data
+        self.frame_lock = threading.Lock()
+        self.black_frame = np.zeros((640, 640, 3), dtype=np.uint8)
+        _, buffer = cv2.imencode('.jpg', self.black_frame)
+        self.latest_frame_bytes = buffer.tobytes()
+        self.latest_frame_ts = time.time()
 
-from reachy_mini.utils import create_head_pose
+        # Audio Stream Data
+        self.audio_queue = queue.Queue()
 
-robot_ws = None  # type: Optional[WebSocket]
-robot_loop = None  # type: Optional[asyncio.AbstractEventLoop]
+    def set_robot_connection(self, ws: WebSocket, loop: asyncio.AbstractEventLoop):
+        self.robot_ws = ws
+        self.robot_loop = loop
 
-_black_frame = np.zeros((640, 640, 3), dtype=np.uint8)
-_, _buffer = cv2.imencode('.jpg', _black_frame)
+    def update_frame(self, frame_bytes: bytes):
+        with self.frame_lock:
+            self.latest_frame_bytes = frame_bytes
+            self.latest_frame_ts = time.time()
 
-frame_lock = threading.Lock()
-latest_frame_data = {
-    "bytes": _buffer.tobytes(),
-    "timestamp": time.time()
-}
+    def push_audio(self, audio_bytes: bytes):
+        """
+        Pushes raw audio bytes. 
+
+        If the queue is full (meaning we are lagging), throw away the OLDEST audio.
+        """
+        # Limit queue to ~0.5 seconds of audio (approx 5-6 chunks of 4096 bytes)
+        MAX_QUEUE_SIZE = 6 
+        
+        while self.audio_queue.qsize() >= MAX_QUEUE_SIZE:
+            try:
+                # Drop the oldest chunk to make room for the new one (Keep 'Now', drop 'Past')
+                print("Dropping oldest audio, queue size is", self.audio_queue.qsize())
+                self.audio_queue.get_nowait()
+            except queue.Empty:
+                pass
+        
+        self.audio_queue.put(audio_bytes)
+
+    def get_connection_status(self) -> str:
+        if self.robot_ws:
+            return "✅ Robot Connected"
+        return "🔴 Waiting for Robot..."
 
+state = GlobalState()
+
+
+# --- 2. Data Models & Presets ---
 @dataclass
 class Movement:
     name: str
@@ -38,675 +120,386 @@ class Movement:
     roll: float = 0
     pitch: float = 0
     yaw: float = 0
+    body_yaw: float = 0
     left_antenna: Optional[float] = None
     right_antenna: Optional[float] = None
     duration: float = 1.0
 
-
-# Preset movements library
-PRESET_MOVEMENTS = {
-    "Home": Movement("Home", 0, 0, 0, 0, 0, 0, 0, 0),
-    "Look Left": Movement("Look Left", 0, 0, 0, 0, 0, -30),
-    "Look Right": Movement("Look Right", 0, 0, 0, 0, 0, 30),
-    "Look Up": Movement("Look Up", 0, 0, 0, 0, -20, 0),
-    "Look Down": Movement("Look Down", 0, 0, 0, 0, 15, 0),
-    "Tilt Left": Movement("Tilt Left", 0, 0, 0, -20, 0, 0),
-    "Tilt Right": Movement("Tilt Right", 0, 0, 0, 20, 0, 0),
-    "Curious": Movement("Curious", 10, 0, 10, 15, -10, -15, 45, -45),
-    "Excited": Movement("Excited", 0, 0, 20, 0, -15, 0, 90, 90),
-    "Shy": Movement("Shy", -10, 0, -10, 10, 10, 20, -30, 30),
+PRESETS = {
+    "Home": Movement("Home", 0, 0, 0, 0, 0, 0, 0, 0, 0),
+    "Look Left": Movement("Look Left", 0, 0, 0, 0, 0, 30, 1, 0, 0),
+    "Look Right": Movement("Look Right", 0, 0, 0, 0, 0, -30, -1, 0, 0),
+    "Look Up": Movement("Look Up", 0, 0, 0, 0, -20, 0, 0, 0, 0),
+    "Look Down": Movement("Look Down", 0, 0, 0, 0, 15, 0, 0, 0, 0),
+    "Curious": Movement("Curious", 10, 0, 10, 15, -10, -15, 0,  45, -45),
+    "Excited": Movement("Excited", 0, 0, 20, 0, -15, 0, 0, 90, 90),
+    "Shy": Movement("Shy", -10, 0, -10, 10, 10, 20, 0, -30, 30),
 }
 
-# Preset sequences
-PRESET_SEQUENCES = {
+SEQUENCES = {
     "Wave": ["Home", "Look Left", "Look Right", "Look Left", "Look Right", "Home"],
     "Nod": ["Home", "Look Down", "Look Up", "Look Down", "Home"],
-    "Excited Dance": ["Home", "Excited", "Tilt Left", "Tilt Right", "Tilt Left", "Home"],
-    "Look Around": ["Home", "Look Left", "Look Up", "Look Right", "Look Down", "Home"],
-    "Curious Peek": ["Home", "Curious", "Look Right", "Look Left", "Home"],
+    "Excited Dance": ["Home", "Excited", "Look Left", "Look Right", "Home"],
 }
 
 
-class ReachyController:
+# --- 3. Controller Logic ---
+class MovementManager:
     def __init__(self):
-        self.daemon_process = None
-        self.running = False
-        
-        self.movement_queue: List[Movement] = []
+        self.queue: List[Movement] = []
         self.is_playing = False
+        self.auto_play = True
         self.playback_speed = 1.0
-        self.play_thread = None
-        self.auto_play = True  # Auto-play mode enabled by default
-        
-        # Connection settings
-        self.connection_mode = "external"  # "local" or "external"
-        self.external_ip = "192.168.1.100"  # Default external IP
-        
-    def start_daemon(self):
-        """Start the Reachy Mini daemon (local mode only)"""
-        try:
-            if self.connection_mode != "local":
-                return "ℹ️ External mode: daemon should be running on remote host"
-            
-            if self.daemon_process is not None:
-                return "⚠️ Daemon already running"
-            
-            python_cmd = "mjpython" if sys.platform == "darwin" else sys.executable
-            self.daemon_process = subprocess.Popen(
-                [python_cmd, "-m", "reachy_mini.daemon.app.main", "--sim",
-                 "--scene", "minimal", "--headless", "--websocket-uri", "ws://localhost:7860", 
-                 "--log-file", "daemon.log", "--deactivate-audio"],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                text=True
-            )
-            
-            frame_received = False  # TODO: check if video is streaming
-            
-            if self.daemon_process.poll() is not None:
-                return "❌ Daemon failed to start"
-            
-            return "✅ Daemon started" if frame_received else "⚠️ Daemon started (no video)"
-        except Exception as e:
-            return f"❌ Daemon error: {str(e)}"
-   
-    def generate_mjpeg_stream(self):
-            global latest_frame_data, frame_lock
-            last_timestamp = 0.0
-            
-            while True:
-                # 1. Check if frame has changed
-                with frame_lock:
-                    current_bytes = latest_frame_data["bytes"]
-                    current_timestamp = latest_frame_data["timestamp"]
-                
-                # 2. Only yield if this is a new frame
-                if current_timestamp > last_timestamp:
-                    last_timestamp = current_timestamp
-                    if current_bytes is not None:
-                        yield (b'--frame\r\n'
-                            b'Content-Type: image/jpeg\r\n\r\n' + current_bytes + b'\r\n')
-                else:
-                    # If no new frame, sleep a bit longer to save CPU
-                    time.sleep(0.02)
-                    continue
-                
-                # Cap FPS slightly to prevent saturation
-                time.sleep(0.02)
+        self.play_thread: Optional[threading.Thread] = None
 
-    def auto_start(self):
-        """Auto-start daemon and robot connection"""
-        status_msgs = []
-        
-        # Start daemon (only for local mode)
-        if self.connection_mode == "local":
-            msg = self.start_daemon()
-            status_msgs.append(msg)
-            yield "\n".join(status_msgs)
-        else:
-            # For external mode, just start frame listener
-            self.start_frame_listener()
-            status_msgs.append(f"ℹ️ External mode: connecting to {self.external_ip}")
-            yield "\n".join(status_msgs)
-        
-    def restart_system(self):
-        """Restart daemon and robot connection"""
-        yield from self.stop_all()
-        yield "🔄 Restarting..."
-        yield from self.auto_start()
-    
-    def stop_all(self):
-        """Stop everything"""
-        self.is_playing = False
-        
-        if self.daemon_process:
-            self.daemon_process.terminate()
-            self.daemon_process.wait(timeout=5)
-            self.daemon_process = None
-        
-        self.running = False
-        
-        return "✅ Stopped"
-    
-    def add_to_queue(self, movement_name, x, y, z, roll, pitch, yaw, 
-                     left_ant, right_ant, duration):
-        """Add a movement to the queue"""
-        movement = Movement(
-            name=movement_name or f"Custom {len(self.movement_queue) + 1}",
-            x=x, y=y, z=z,
-            roll=roll, pitch=pitch, yaw=yaw,
-            left_antenna=left_ant,
-            right_antenna=right_ant,
-            duration=duration
-        )
-        self.movement_queue.append(movement)
-        
-        # Auto-play if enabled and not already playing
+    def add_movement(self, mov: Movement):
+        self.queue.append(mov)
         if self.auto_play and not self.is_playing:
-            self._start_auto_play()
-        
-        return self.format_queue(), f"✅ Added: {movement.name}"
-    
-    def add_preset(self, preset_name):
-        """Add a preset movement to queue"""
-        if preset_name not in PRESET_MOVEMENTS:
-            return self.format_queue(), f"❌ Unknown preset: {preset_name}"
-        
-        self.movement_queue.append(PRESET_MOVEMENTS[preset_name])
-        
-        # Auto-play if enabled and not already playing
-        if self.auto_play and not self.is_playing:
-            self._start_auto_play()
-        
-        return self.format_queue(), f"✅ Added: {preset_name}"
-    
-    def add_sequence(self, sequence_name):
-        """Add a preset sequence to queue"""
-        if sequence_name not in PRESET_SEQUENCES:
-            return self.format_queue(), "❌ Unknown sequence"
-        
-        for preset_name in PRESET_SEQUENCES[sequence_name]:
-            self.movement_queue.append(PRESET_MOVEMENTS[preset_name])
-        
-        # Auto-play if enabled and not already playing
-        if self.auto_play and not self.is_playing:
-            self._start_auto_play()
-        
-        return self.format_queue(), f"✅ Added sequence: {sequence_name}"
-    
+            self.play_queue()
+        return self.get_queue_text(), f"✅ Added {mov.name}"
+
+    def add_preset(self, name: str):
+        if name in PRESETS:
+            return self.add_movement(PRESETS[name])
+        return self.get_queue_text(), "❌ Unknown Preset"
+
+    def add_sequence(self, name: str):
+        if name in SEQUENCES:
+            for move_name in SEQUENCES[name]:
+                self.queue.append(PRESETS[move_name])
+            if self.auto_play and not self.is_playing:
+                self.play_queue()
+            return self.get_queue_text(), f"✅ Added Sequence: {name}"
+        return self.get_queue_text(), "❌ Unknown Sequence"
+
     def clear_queue(self):
-        """Clear the movement queue"""
-        self.movement_queue.clear()
-        self.is_playing = False  # Stop playback when clearing
-        return self.format_queue(), "🗑️ Queue cleared"
-    
+        self.queue.clear()
+        self.is_playing = False
+        return self.get_queue_text(), "🗑️ Queue Cleared"
+
     def remove_last(self):
-        """Remove last movement from queue"""
-        if self.movement_queue:
-            removed = self.movement_queue.pop()
-            return self.format_queue(), f"🗑️ Removed: {removed.name}"
-        return self.format_queue(), "⚠️ Queue is empty"
-    
-    def format_queue(self):
-        """Format queue for display"""
-        if not self.movement_queue:
-            return "📋 Queue is empty\n\nAdd movements using presets or custom controls"
-        
-        lines = ["📋 Movement Queue:\n"]
-        total_duration = 0
+        if self.queue:
+            self.queue.pop()
+        return self.get_queue_text(), "🗑️ Removed Last"
+
+    def get_queue_text(self):
+        if not self.queue:
+            return "📋 Queue Empty"
         
-        for i, mov in enumerate(self.movement_queue, 1):
-            total_duration += mov.duration
-            emoji = "▶️" if i == 1 else "⏸️"
-            
-            # Format head position
-            head_str = f"Head: x={mov.x:.0f} y={mov.y:.0f} z={mov.z:.0f} r={mov.roll:.0f}° p={mov.pitch:.0f}° y={mov.yaw:.0f}°"
-            
-            # Format antennas if present
-            ant_str = ""
-            if mov.left_antenna is not None and mov.right_antenna is not None:
-                ant_str = f"\n   Antennas: L={mov.left_antenna:.0f}° R={mov.right_antenna:.0f}°"
-            
-            lines.append(
-                f"{emoji} {i}. {mov.name} ({mov.duration}s)\n"
-                f"   {head_str}{ant_str}"
-            )
+        lines = ["📋 Current Queue:"]
+        for i, m in enumerate(self.queue, 1):
+            indicator = "▶️" if (i==1 and self.is_playing) else "  "
+            lines.append(f"{indicator} {i}. {m.name} ({m.duration}s)")
         
-        lines.append(f"\n⏱️ Total duration: {total_duration:.1f}s")
-        lines.append(f"{'🔄 Auto-play: ON' if self.auto_play else '⏸️ Auto-play: OFF'}")
         return "\n".join(lines)
-    
-    def play_queue(self, speed):
-        """Execute the movement queue"""
-        if not self.movement_queue:
-            return self.format_queue(), "⚠️ Queue is empty"
-        
-        if self.is_playing:
-            return self.format_queue(), "⚠️ Already playing"
-        
-        self.playback_speed = speed
+
+    def play_queue(self, speed=None):
+        if speed: self.playback_speed = speed
+        if self.is_playing: return self.get_queue_text(), "⚠️ Already Playing"
+        if not self.queue: return self.get_queue_text(), "⚠️ Queue Empty"
+
         self.is_playing = True
-        self.play_thread = threading.Thread(target=self._play_loop, daemon=True)
+        self.play_thread = threading.Thread(target=self._worker, daemon=True)
         self.play_thread.start()
-        
-        return self.format_queue(), f"▶️ Playing at {speed}x speed..."
-    
-    def _play_loop(self):
-        """Background thread to execute movements"""
+        return self.get_queue_text(), "▶️ Playing..."
+
+    def stop_playback(self):
+        self.is_playing = False
+        if self.play_thread:
+            self.play_thread.join(timeout=1.0)
+        return self.get_queue_text(), "⏹️ Stopped"
+
+    def _worker(self):
+        """Background thread that processes the queue."""
+        idx = 0
         try:
-            current_index = 0
-            while self.is_playing:
-                # Check if there are movements to play
-                if current_index < len(self.movement_queue):
-                    print("[Space] Playing movement:", current_index)
-                    movement = self.movement_queue[current_index]
-                    
-                    # Create pose
-                    pose = create_head_pose(
-                        x=movement.x, y=movement.y, z=movement.z,
-                        roll=movement.roll, pitch=movement.pitch, yaw=movement.yaw,
-                        degrees=True, mm=True
+            while self.is_playing and idx < len(self.queue):
+                move = self.queue[idx]
+                
+                # 1. Build Payload
+                pose = create_head_pose(
+                    x=move.x, y=move.y, z=move.z,
+                    roll=move.roll, pitch=move.pitch, yaw=move.yaw,
+                    degrees=True, mm=True
+                )
+                
+                payload = {
+                    "type": "movement",
+                    "movement": {
+                        "head": pose.tolist(),
+                        "body_yaw": move.body_yaw,
+                        "duration": move.duration / self.playback_speed
+                    }
+                }
+
+                # Add antennas if specified
+                if move.left_antenna is not None and move.right_antenna is not None:
+                    payload["movement"]["antennas"] = [
+                        np.deg2rad(move.right_antenna), 
+                        np.deg2rad(move.left_antenna)
+                    ]
+
+                # 2. Send to Robot (Async safe)
+                if state.robot_ws and state.robot_loop:
+                    asyncio.run_coroutine_threadsafe(
+                        state.robot_ws.send_json(payload), 
+                        state.robot_loop
                     )
-                    
-                    # Adjust duration by playback speed
-                    actual_duration = movement.duration / self.playback_speed
-                    
-                    
-                    if movement.left_antenna is not None and movement.right_antenna is not None:
-                        send_robot_command({
-                            "type": "movement",
-                            "movement": {
-                                "head": pose.tolist(),
-                                "antennas": [np.deg2rad(movement.right_antenna), np.deg2rad(movement.left_antenna)],
-                                "duration": actual_duration
-                            }
-                        })
-                    else:
-                        send_robot_command({
-                            "type": "movement",
-                            "movement": {
-                                "head": pose.tolist(),
-                                "duration": actual_duration
-                            }
-                        })
-                    
-                    current_index += 1
                 else:
-                    # No more movements, stop if not in auto-play mode
-                    if not self.auto_play:
-                        break
-                    # In auto-play mode, wait for new movements
-                    time.sleep(0.1)
-            
+                    print("⚠️ Robot not connected, skipping command.")
+
+                # 3. Wait for move to finish (blocking thread, not async loop)
+                time.sleep(move.duration / self.playback_speed)
+                idx += 1
+
+            # Loop finished
+            if not self.auto_play:
+                self.is_playing = False
+            else:
+                # In auto-play, we stay "playing" but wait for new items
+                self.is_playing = False 
+
         except Exception as e:
-            print(f"Error during playback: {e}")
-        finally:
+            print(f"Playback Error: {e}")
             self.is_playing = False
-    
-    def _start_auto_play(self):
-        """Start auto-play mode"""
-        if not self.is_playing:
-            self.is_playing = True
-            self.play_thread = threading.Thread(target=self._play_loop, daemon=True)
-            self.play_thread.start()
-    
-    def toggle_auto_play(self, enabled):
-        """Toggle auto-play mode"""
-        self.auto_play = enabled
-        
-        if self.auto_play and self.movement_queue and not self.is_playing:
-            self._start_auto_play()
-        
-        return self.format_queue(), f"{'🔄 Auto-play enabled' if enabled else '⏸️ Auto-play disabled'}"
-    
-    def update_speed(self, speed):
-        """Update playback speed in real-time"""
-        self.playback_speed = speed
-        return f"⚡ Speed: {speed}x"
-    
-    def stop_playback(self):
-        """Stop current playback"""
-        self.is_playing = False
-        if self.play_thread:
-            self.play_thread.join(timeout=2)
-        
-        # If auto-play is still enabled, inform user
-        msg = "⏹️ Stopped"
-        if self.auto_play:
-            msg += " (auto-play still enabled)"
-        
-        return self.format_queue(), msg
-    
-    def set_connection_mode(self, mode, external_ip=None):
-        """Set connection mode and optionally external IP"""
-        self.connection_mode = mode
-        print(f"External IP: {external_ip}")
-        if external_ip:
-            self.external_ip = external_ip
-        
-        yield f"🔄 Connection mode set to: {mode}\n🔄 Restarting..."
-        yield from self.restart_system()
 
+manager = MovementManager()
 
-# Create manager
-manager = ReachyController()
 
-# Build Gradio interface with improved layout
-with gr.Blocks(title="Reachy Controller", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🤖 Reachy Mini Controller")
-    gr.Markdown("Create fun movement sequences for your robot!")
+# --- 4. FastAPI & WebSocket Logic ---
+app = FastAPI()
+
+@app.websocket("/robot")
+async def robot_endpoint(ws: WebSocket):
+    """Endpoint for the Robot to connect to (Control Channel)."""
+    await ws.accept()
+    state.set_robot_connection(ws, asyncio.get_running_loop())
+    print("[System] Robot Connected!")
     
-    with gr.Row():
-        # Left panel - Controls (narrower)
-        with gr.Column(scale=2):
-            # Connection mode section
-            gr.Markdown("### 🔌 Connection Settings")
-            connection_mode = gr.Radio(
-                choices=["local", "external"],
-                value="external",
-                label="Connection Mode",
-                info="Local: Launch daemon locally | External: Connect to remote daemon"
-            )
-            external_ip_input = gr.Textbox(
-                label="External IP or Hostname",
-                value="192.168.1.100",
-                placeholder="192.168.1.100 or 8.tcp.ngrok.io:18951",
-                visible=True,
-                info="IP/hostname of remote daemon. Supports custom ports (e.g., ngrok.io:18951)"
-            )
-            
-            def toggle_external_ip_visibility(mode):
-                return gr.update(visible=(mode == "external"))
-            
-            connection_mode.change(
-                fn=toggle_external_ip_visibility,
-                inputs=[connection_mode],
-                outputs=[external_ip_input]
-            )
-            
-            apply_connection_btn = gr.Button("✅ Apply Connection Settings", variant="primary", size="sm")
-            
-            # Status section
-            gr.Markdown("### 📡 System Status")
-            status = gr.Textbox(
-                label="Status",
-                lines=3,
-                interactive=False,
-                value="🔄 Initializing system..."
-            )
-            restart_btn = gr.Button("🔄 Restart System", variant="secondary", size="sm")
+    try:
+        # Heartbeat loop
+        while True:
+            # We wait for messages, but mostly we just hold the connection open
+            # and send commands via the state.robot_ws handle.
+            msg = await ws.receive() 
+            if msg["type"] == "websocket.disconnect":
+                break
+    except (WebSocketDisconnect, Exception):
+        print("[System] Robot Disconnected")
+    finally:
+        state.robot_ws = None
+
+@app.websocket("/video_stream")
+async def stream_endpoint(ws: WebSocket):
+    """Endpoint for Robot/Sim to send video frames."""
+    await ws.accept()
+    try:
+        while True:
+            msg = await ws.receive()
+            if "bytes" in msg and msg["bytes"]:
+                state.update_frame(msg["bytes"])
+    except Exception:
+        pass
+
+@app.websocket("/audio_stream")
+async def audio_endpoint(ws: WebSocket):
+    """Endpoint for Robot/Sim to send raw Audio bytes (int16)."""
+    await ws.accept()
+    print("[Audio] Stream Connected")
+    try:
+        while True:
+            # Receive bytes directly
+            data = await ws.receive()
+            if data.get('type') == 'websocket.receive' and data.get('bytes'):
+                state.push_audio(data.get('bytes'))
+            elif data.get('type') == 'websocket.receive' and data.get('text') == "ping":
+                print("[Audio] Received ping")
+            elif data.get('type') == 'websocket.disconnect':
+                print("[Audio] Disconnected")
+                break
+            else:
+                print(f"[Audio] Received unknown message: {data}")
+    except Exception as e:
+        print(f"[Audio] Disconnected: {e}")
 
-            test_cmd_btn = gr.Button("📡 Send test command", size="sm")
+# --- 5. Gradio Interface ---
 
-            def send_test_command():
-                send_robot_command({"type": "test", "message": "hello from space"})
-                return "Sent test command to robot"
+def webrtc_audio_generator():
+    """
+    Generator for FastRTC.
+    """
+    # Clear old data to start fresh
+    with state.audio_queue.mutex:
+        state.audio_queue.queue.clear()
 
+    # OPTIMIZATION: Reduce target samples. 
+    # 4096 samples @ 16kHz is 256ms of latency built-in! 
+    # Try 1024 (64ms) or 512 (32ms) for lower latency.
+    TARGET_SAMPLES = 1024 
+    byte_buffer = bytearray()
 
-            test_cmd_btn.click(
-                fn=send_test_command,
-                outputs=[status],  # reuse your status textbox
-            )
-            
-            gr.Markdown("### 🎮 Playback Controls")
-            
-            auto_play_toggle = gr.Checkbox(
-                label="🔄 Auto-play",
-                value=True,
-                info="Execute movements automatically when added"
-            )
-            
-            speed_slider = gr.Slider(
-                0.25, 3.0, 1.0,
-                label="⚡ Speed Multiplier",
-                info="Adjust playback speed"
-            )
-            
-            with gr.Row():
-                play_btn = gr.Button("▶️ Play All", variant="primary", scale=2)
-                stop_play_btn = gr.Button("⏹️ Stop", scale=1)
-            
-            with gr.Row():
-                clear_btn = gr.Button("🗑️ Clear All")
-                remove_btn = gr.Button("↶ Remove Last")
-            
-            # Queue display
-            queue_display = gr.Textbox(
-                label="📋 Movement Queue",
-                lines=20,
-                interactive=False,
-                value=manager.format_queue()
-            )
-        
-        # Right panel - Simulation view (larger and more square)
-        with gr.Column(scale=3):
-            # sim_view = gr.Image(
-            #     label="🎬 Robot Simulation",
-            #     type="numpy",
-            #     height=1080,
-            #     width=1080,
-            #     show_label=True,
-            #     streaming=True
-            # )
-            html_code = """
-                <html>
-                <body>
-                    <img src="/video_feed" style="width: 100%; max-width: 1080px; border-radius: 8px;">
-                </body>
-                </html>
-                """
-            sim_view = gr.HTML(value=html_code, label="🎬 Robot Simulation")
+    while True:
+        try:
+            # Wait up to 1 second for data. If no data, loop again.
+            # Do NOT use a short timeout combined with silence generation.
+            chunk_bytes = state.audio_queue.get(timeout=1.0)
+            if chunk_bytes:
+                byte_buffer.extend(chunk_bytes)
+        except queue.Empty:
+            # If we really have no data for a long time, just continue waiting.
+            # Do NOT yield silence here.
+            continue
+
+        # Only yield when we have enough data
+        while len(byte_buffer) >= TARGET_SAMPLES * 2:  # int16 = 2 bytes
+            read_size = TARGET_SAMPLES * 2
+            out_bytes = byte_buffer[:read_size]
+            byte_buffer = byte_buffer[read_size:]
+
+            audio_int16 = np.frombuffer(out_bytes, dtype=np.int16)
+            audio_int16 = audio_int16.reshape(1, -1)
+
+            yield (AUDIO_SAMPLE_RATE, audio_int16)
+
+def webrtc_video_generator():
+    """
+    Generator for FastRTC WebRTC (mode='receive', modality='video').
+
+    It reads JPEG bytes from state.latest_frame_bytes, decodes them with OpenCV,
+    and yields HxWx3 uint8 frames as expected by FastRTC.
+    """
+    last_ts = 0.0
+    frame = state.black_frame.copy()
+
+    while True:
+        with state.frame_lock:
+            ts = state.latest_frame_ts
+            frame_bytes = state.latest_frame_bytes
+
+        if ts > last_ts and frame_bytes:
+            last_ts = ts
+            np_bytes = np.frombuffer(frame_bytes, dtype=np.uint8)
+            frame = cv2.imdecode(np_bytes, cv2.IMREAD_COLOR)
+            if frame is None:
+                frame = state.black_frame.copy()
+        # Shape (H, W, 3), dtype uint8
+        yield frame
 
+with gr.Blocks(title="Reachy Controller", theme=gr.themes.Soft()) as demo:
+    
+    gr.Markdown("## 🤖 Reachy Mini Controller")
     
-    # Movement builder section below
     with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 🎨 Quick Presets")
-            
-            with gr.Row():
-                preset_btns = []
-                for preset in list(PRESET_MOVEMENTS.keys())[:5]:
-                    btn = gr.Button(preset, size="sm")
-                    preset_btns.append((btn, preset))
+        # --- LEFT COLUMN: Controls ---
+        with gr.Column(scale=1):
+            status_box = gr.Textbox(label="System Status", value=state.get_connection_status, every=2)
             
-            with gr.Row():
-                for preset in list(PRESET_MOVEMENTS.keys())[5:]:
-                    btn = gr.Button(preset, size="sm")
-                    preset_btns.append((btn, preset))
-        
-        with gr.Column():
-            gr.Markdown("### 🎬 Sequences")
-            with gr.Row():
-                sequence_dropdown = gr.Dropdown(
-                    choices=list(PRESET_SEQUENCES.keys()),
-                    label="Select Sequence",
-                    value=None,
-                    scale=3
+            with gr.Group():
+                gr.Markdown("### 🎧 Audio Listen")
+
+                # Start button for the WebRTC stream
+                listen_btn = gr.Button("🔊 Start Listening", variant="secondary")
+
+                # FastRTC WebRTC component in receive mode, audio only
+                robot_audio = WebRTC(
+                    label="Robot Audio",
+                    modality="audio",
+                    mode="receive",
+                    # Optional niceties:
+                    # icon="phone-solid.svg",
+                    # icon_button_color="black",
+                    # pulse_color="black",
                 )
-                add_seq_btn = gr.Button("➕ Add", scale=1)
-    
-    # Custom movement controls in accordion
-    with gr.Accordion("🎯 Custom Movement Builder", open=False):
-        custom_name = gr.Textbox(label="Movement Name", placeholder="My Move")
-        
-        with gr.Row():
-            x = gr.Slider(-50, 50, 0, label="X (mm)", step=5)
-            y = gr.Slider(-50, 50, 0, label="Y (mm)", step=5)
-            z = gr.Slider(-20, 50, 0, label="Z (mm)", step=5)
-        
-        with gr.Row():
-            roll = gr.Slider(-30, 30, 0, label="Roll (°)", step=5)
-            pitch = gr.Slider(-30, 30, 0, label="Pitch (°)", step=5)
-            yaw = gr.Slider(-45, 45, 0, label="Yaw (°)", step=5)
-        
-        with gr.Row():
-            left_ant = gr.Slider(-180, 180, 0, label="Left Antenna (°)", step=15)
-            right_ant = gr.Slider(-180, 180, 0, label="Right Antenna (°)", step=15)
-        
-        duration = gr.Slider(0.3, 3.0, 1.0, label="Duration (s)", step=0.1)
-        
-        add_custom_btn = gr.Button("➕ Add to Queue", variant="primary")
-    
-    # Connect events - Connection settings
-    apply_connection_btn.click(
-        fn=manager.set_connection_mode,
-        inputs=[connection_mode, external_ip_input],
-        outputs=[status]
-    )
-    
-    # Connect events - System control
-    restart_btn.click(fn=manager.restart_system, outputs=[status])
-    
-    # Connect events - Playback control
-    auto_play_toggle.change(
-        fn=manager.toggle_auto_play,
-        inputs=[auto_play_toggle],
-        outputs=[queue_display, status]
-    )
-    
-    speed_slider.change(
-        fn=manager.update_speed,
-        inputs=[speed_slider],
-        outputs=[status]
-    )
-    
-    play_btn.click(
-        fn=manager.play_queue,
-        inputs=[speed_slider],
-        outputs=[queue_display, status]
-    )
-    stop_play_btn.click(fn=manager.stop_playback, outputs=[queue_display, status])
-    clear_btn.click(fn=manager.clear_queue, outputs=[queue_display, status])
-    remove_btn.click(fn=manager.remove_last, outputs=[queue_display, status])
-    
-    # Connect preset buttons
-    for btn, preset_name in preset_btns:
-        btn.click(
-            fn=lambda p=preset_name: manager.add_preset(p),
-            outputs=[queue_display, status]
-        )
-    
-    # Connect sequence dropdown
-    add_seq_btn.click(
-        fn=manager.add_sequence,
-        inputs=[sequence_dropdown],
-        outputs=[queue_display, status]
-    )
-    
-    # Connect custom movement
-    add_custom_btn.click(
-        fn=manager.add_to_queue,
-        inputs=[custom_name, x, y, z, roll, pitch, yaw, left_ant, right_ant, duration],
-        outputs=[queue_display, status]
-    )
-app = FastAPI()
 
-def send_robot_command(cmd: dict):
-    global robot_ws, robot_loop
-    send_start_time = time.perf_counter()
-    print("[Space] Sending command to robot:", cmd)
-
-    if robot_ws is None or robot_loop is None:
-        print("[Space] Cannot send command, robot not connected")
-        return
+                # Wire generator to WebRTC component
+                robot_audio.stream(
+                    fn=lambda: webrtc_audio_generator(),
+                    inputs=None,
+                    outputs=[robot_audio],
+                    trigger=listen_btn.click,
+                )
+                
+                
+            with gr.Group():
+                gr.Markdown("### 🎮 Playback")
+                auto_play = gr.Checkbox(label="Auto-play", value=True)
+                speed = gr.Slider(0.5, 2.0, 1.0, label="Speed")
+                
+                with gr.Row():
+                    play_btn = gr.Button("▶️ Play", variant="primary")
+                    stop_btn = gr.Button("⏹️ Stop")
+                
+                with gr.Row():
+                    clear_btn = gr.Button("🗑️ Clear")
+                    undo_btn = gr.Button("↶ Undo")
 
-    async def _send():
-        global robot_ws
-        try:
-            send_ws_start = time.perf_counter()
-            await robot_ws.send_json(cmd)  # or send_text(json.dumps(cmd)) if you prefer
-            send_ws_time = time.perf_counter() - send_ws_start
-            print(f"[WebSocket /robot] Send time: {send_ws_time*1000:.2f}ms")
-        except Exception as e:
-            print("[Space] Error sending command to robot:", e)
-            robot_ws = None
+            queue_display = gr.Textbox(label="Queue", value=manager.get_queue_text, lines=10)
 
-    # Schedule coroutine on the main event loop from this worker thread
-    fut = asyncio.run_coroutine_threadsafe(_send(), robot_loop)
+        # --- RIGHT COLUMN: View ---
+        with gr.Column(scale=2):
+            robot_video = WebRTC(
+                label="Robot Video",
+                modality="video",
+                mode="receive",
+            )
+            robot_video.stream(
+                fn=lambda: webrtc_video_generator(),
+                inputs=None,
+                outputs=[robot_video],
+                trigger=listen_btn.click,
+            )
 
-    # Optionally wait for completion (and surface errors)
-    try:
-        fut.result(timeout=2)
-        total_send_time = time.perf_counter() - send_start_time
-        print(f"[WebSocket /robot] Total send time (including scheduling): {total_send_time*1000:.2f}ms")
-    except Exception as e:
-        print("[Space] Command failed:", e)
-        robot_ws = None
-   
-@app.websocket("/mujoco_stream")
-async def robot_socket(ws: WebSocket):
-    global latest_frame_data, frame_lock
-    await ws.accept()
-    print("[Space] MuJoCo stream connected")
-    
-    try:
-        while True:
-            msg = await ws.receive()
-            
-            if msg["type"] == "websocket.receive":
-                # 1. Safe retrieval using .get() to avoid KeyError
-                payload_bytes = msg.get("bytes")
-                payload_text = msg.get("text")
-
-                # 2. Handle Binary Frame
-                if payload_bytes is not None:
-                    with frame_lock:
-                        latest_frame_data["bytes"] = payload_bytes
-                        latest_frame_data["timestamp"] = time.time()
-                        
-                # 3. Handle Text Message (The Ping)
-                elif payload_text is not None:
-                    try:
-                        data = json.loads(payload_text)
-                        if data.get("type") == "ping":
-                            # Optional: Print debug only if you need to verify it's working
-                            print("[Space] Frame keep-alive received")
-                            pass
-                    except json.JSONDecodeError:
-                        print(f"[Space] Received invalid JSON: {payload_text}")
-                        
-            elif msg["type"] == "websocket.disconnect":
-                print("[Space] MuJoCo stream disconnected")
-                break
-                
-    except Exception as e:
-        print(f"[Space] MuJoCo stream disconnected error: {e}")
+            # --- Movement Builders ---
+            with gr.Tabs():
+                with gr.Tab("✨ Presets & Sequences"):
+                    gr.Markdown("### Quick Actions")
+                    with gr.Row(variant="panel"):
+                        for name in PRESETS:
+                            btn = gr.Button(name, size="sm")
+                            btn.click(manager.add_preset, inputs=[gr.State(name)], outputs=[queue_display, status_box])
+                    
+                    gr.Markdown("### Sequences")
+                    with gr.Row():
+                        for seq in SEQUENCES:
+                            btn = gr.Button(f"🎬 {seq}", size="sm")
+                            btn.click(manager.add_sequence, inputs=[gr.State(seq)], outputs=[queue_display, status_box])
+
+                with gr.Tab("🛠️ Custom Move"):
+                    with gr.Row():
+                        c_x = gr.Slider(-50, 50, 0, label="X")
+                        c_y = gr.Slider(-50, 50, 0, label="Y")
+                        c_z = gr.Slider(-20, 50, 0, label="Z")
+                    with gr.Row():
+                        c_r = gr.Slider(-30, 30, 0, label="Roll")
+                        c_p = gr.Slider(-30, 30, 0, label="Pitch")
+                        c_y_aw = gr.Slider(-45, 45, 0, label="Yaw")
+                    with gr.Row():
+                        c_la = gr.Slider(-180, 180, 0, label="Left Ant")
+                        c_ra = gr.Slider(-180, 180, 0, label="Right Ant")
+                    
+                    c_dur = gr.Slider(0.1, 5.0, 1.0, label="Duration")
+                    c_add = gr.Button("➕ Add Custom Move", variant="primary")
 
-@app.get("/video_feed")
-def video_feed():
-    return StreamingResponse(
-        manager.generate_mjpeg_stream(), 
-        media_type="multipart/x-mixed-replace; boundary=frame"
-    )
+                    def _add_custom(x,y,z,r,p,yw,la,ra,d):
+                        m = Movement("Custom", x,y,z,r,p,yw,la,ra,d)
+                        return manager.add_movement(m)
 
-@app.websocket("/robot")
-async def robot_socket(ws: WebSocket):
-    global robot_ws, robot_loop
-    await ws.accept()
-    robot_ws = ws
-    robot_loop = asyncio.get_running_loop()
-    print("[Space] Robot connected")
-    
-    # --- HEARTBEAT MECHANISM ---
-    # Define a keep-alive task that runs in the background
-    async def keep_alive():
-        try:
-            while True:
-                await asyncio.sleep(30) # Send ping every 30s
-                if robot_ws:
-                    # Send a lightweight ping
-                    await robot_ws.send_json({"type": "ping"}) 
-        except Exception as e:
-            print(f"[Space] Heartbeat stopped: {e}")
+                    c_add.click(_add_custom, 
+                                inputs=[c_x, c_y, c_z, c_r, c_p, c_y_aw, c_la, c_ra, c_dur], 
+                                outputs=[queue_display, status_box])
 
-    # Start the heartbeat
-    heartbeat_task = asyncio.create_task(keep_alive())
+    # --- Event Wiring ---
+    auto_play.change(lambda x: setattr(manager, 'auto_play', x), inputs=[auto_play])
+    play_btn.click(manager.play_queue, inputs=[speed], outputs=[queue_display, status_box])
+    stop_btn.click(manager.stop_playback, outputs=[queue_display, status_box])
+    clear_btn.click(manager.clear_queue, outputs=[queue_display, status_box])
+    undo_btn.click(manager.remove_last, outputs=[queue_display, status_box])
 
-    try:
-        while True:
-            msg = await ws.receive()
-            if msg["type"] == "websocket.receive":
-                if msg["text"] is not None:
-                    # Handle robot messages here
-                    pass
-            elif msg["type"] == "websocket.disconnect":
-                print("[Space] Robot disconnected")
-                break
-    except Exception as e:
-        print("[Space] Robot disconnected (Error)", e)
-    finally:
-        heartbeat_task.cancel() # Stop the heartbeat when connection dies
-        if robot_ws is ws:
-            robot_ws = None
 
+# --- 6. Mount & Run ---
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(
-            app, 
-            host="0.0.0.0", 
-            port=7860, 
-            proxy_headers=True,      # Tell Uvicorn to trust the proxy's headers
-            forwarded_allow_ips="*"  # Trust headers from any IP (the internal HF proxy)
-        )
+    print("🚀 Server starting on http://0.0.0.0:7860")
+    print("ℹ️  Point your Robot/Sim to: ws://<YOUR_PC_IP>:7860/robot")
+    uvicorn.run(app, host="0.0.0.0", port=7860, proxy_headers=True, forwarded_allow_ips="*")