app.py

"""Reachy Mini Controller - Fun, Queue-Based Control Interface"""

import asyncio
import subprocess
import sys
import threading
import json
import time
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.responses import StreamingResponse

from reachy_mini.utils import create_head_pose

robot_ws = None  # type: Optional[WebSocket]
robot_loop = None  # type: Optional[asyncio.AbstractEventLoop]

_black_frame = np.zeros((640, 640, 3), dtype=np.uint8)
_, _buffer = cv2.imencode('.jpg', _black_frame)

frame_lock = threading.Lock()
latest_frame_data = {
    "bytes": _buffer.tobytes(),
    "timestamp": time.time()
}

@dataclass
class Movement:
    name: str
    x: float = 0
    y: float = 0
    z: float = 0
    roll: float = 0
    pitch: float = 0
    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),
}

# Preset sequences
PRESET_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"],
}


class ReachyController:
    def __init__(self):
        self.daemon_process = None
        self.reachy_mini = None
        self.running = False
        
        self.movement_queue: List[Movement] = []
        self.is_playing = False
        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", "--stream-robot-view", 
                 "--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 stream_frames(self):
        global latest_test_frame
        frame_count = 0
        fps_start_time = time.perf_counter()
        
        while True:
            frame_start_time = time.perf_counter()
            try:
                # latest_test_frame is a JPEG bytestring as produced by cv2.imencode + .tobytes,
                # which matches what video_ws.py sends.
                data = latest_test_frame
                decode_start = time.perf_counter()
                if isinstance(data, bytes):
                    arr = np.frombuffer(data, dtype=np.uint8)
                    frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)  # returns BGR np.ndarray
                    if frame is None:
                        # Decoding failed, show black frame
                        frame = _black_frame
                    else:
                        # Convert BGR->RGB for Gradio and resize to 1080x1080dd
                        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        # frame = cv2.resize(frame, (1080, 1080), interpolation=cv2.INTER_LINEAR)
                else:
                    # Unexpected type, show black frame
                    frame = _black_frame
                decode_time = time.perf_counter() - decode_start
            except Exception as e:
                print("[stream_frames] JPEG decode error:", e)
                frame = _black_frame
                decode_time = 0
            
            frame_count += 1
            
            # Calculate and print FPS every 30 frames (~1 second at 30fps)
            if frame_count % 10 == 0:
                elapsed = time.perf_counter() - fps_start_time
                fps = frame_count / elapsed
                avg_frame_time = elapsed / frame_count
                print(f"[stream_frames] FPS: {fps:.2f} | Avg frame time: {avg_frame_time*1000:.2f}ms | Last decode: {decode_time*1000:.2f}ms")
                # Reset counters for next measurement period
                frame_count = 0
                fps_start_time = time.perf_counter()
            
            yield frame
            time.sleep(0.04)
            
    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)

    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
        
        if self.frame_thread:
            self.frame_thread.join(timeout=2)
            self.frame_thread = None

        if self.receiver:
            self.receiver.close()
            time.sleep(0.5)  # Wait for the socket to close
            self.receiver = None
        
        if self.reachy_mini:
            try:
                self.reachy_mini.__exit__(None, None, None)
            except:
                pass
            self.reachy_mini = None
        
        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
        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}"
    
    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"
    
    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
        
        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.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
        self.is_playing = True
        self.play_thread = threading.Thread(target=self._play_loop, 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"""
        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
                    )
                    
                    # 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)
            
        except Exception as e:
            print(f"Error during playback: {e}")
        finally:
            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()


# 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!")
    
    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")

            test_cmd_btn = gr.Button("📡 Send test command", size="sm")

            def send_test_command():
                send_robot_command({"type": "test", "message": "hello from space"})
                return "Sent test command to robot"


            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")

    
    # 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))
            
            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
                )
                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")
    
    # Stream video
    # demo.load(fn=manager.stream_frames, outputs=sim_view)

    # 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

    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

    # Schedule coroutine on the main event loop from this worker thread
    fut = asyncio.run_coroutine_threadsafe(_send(), robot_loop)

    # 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}")

@app.get("/video_feed")
def video_feed():
    return StreamingResponse(
        manager.generate_mjpeg_stream(), 
        media_type="multipart/x-mixed-replace; boundary=frame"
    )

@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}")

    # Start the heartbeat
    heartbeat_task = asyncio.create_task(keep_alive())

    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

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)
        )