mujoco_server.py

import os
import sys
import time
import threading
import json
import base64
import math
import re
import traceback
import io
import cv2
import numpy as np


def _configure_mujoco_backend() -> None:
    if os.environ.get("NOVA_MUJOCO_GPU", "").lower() in {"1", "true", "yes", "on"}:
        os.environ["MUJOCO_GL"] = os.environ.get("MUJOCO_GL", "egl")
        os.environ["PYOPENGL_PLATFORM"] = os.environ.get("PYOPENGL_PLATFORM", "egl")
    else:
        os.environ["MUJOCO_GL"] = os.environ.get("MUJOCO_GL", "osmesa")
        os.environ["PYOPENGL_PLATFORM"] = os.environ.get("PYOPENGL_PLATFORM", "osmesa")


_configure_mujoco_backend()

import mujoco
from functools import wraps
from pathlib import Path
from typing import Any, Optional
from flask import Flask, Response, render_template_string, request, jsonify, make_response
from flask_sock import Sock

# Add robots directory to path for imports
_nova_sim_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_nova_sim_dir, 'robots', 'g1'))
from g1_env import G1Env
sys.path.pop(0)

# Load .env.local for Nova API configuration
def load_env_file(filepath):
    """Load environment variables from .env file."""
    if not filepath.exists():
        return
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, value = line.split('=', 1)
                os.environ[key.strip()] = value.strip()

env_local_path = Path(_nova_sim_dir) / '.env.local'
load_env_file(env_local_path)

# API prefix for all endpoints
API_PREFIX = '/nova-sim/api/v1'

# Determine if Nova credentials are present (required for default toggle state)
def _nova_credentials_present() -> bool:
    instance_url = os.getenv('NOVA_INSTANCE_URL') or os.getenv('NOVA_API')
    return all([
        instance_url,
        os.getenv('NOVA_ACCESS_TOKEN'),
        os.getenv('NOVA_CONTROLLER_ID'),
        os.getenv('NOVA_MOTION_GROUP_ID'),
    ])

app = Flask(__name__)
app.config['SECRET_KEY'] = 'robotsim-secret'
sock = Sock(app)

# Detect if running in Docker (check for /.dockerenv or cgroup)
def is_running_in_docker():
    if os.path.exists('/.dockerenv'):
        return True
    try:
        with open('/proc/1/cgroup', 'r') as f:
            return 'docker' in f.read()
    except:
        return False

IN_DOCKER = is_running_in_docker()

# Performance settings from environment
# Native defaults: 1280x720 @ 60 FPS (GPU accelerated)
# Docker defaults: 640x360 @ 30 FPS (software rendering)
RENDER_WIDTH = int(os.environ.get('RENDER_WIDTH', 640 if IN_DOCKER else 1280))
RENDER_HEIGHT = int(os.environ.get('RENDER_HEIGHT', 360 if IN_DOCKER else 720))
TARGET_FPS = int(os.environ.get('TARGET_FPS', 30 if IN_DOCKER else 60))
SIM_STEPS_PER_FRAME = int(os.environ.get('SIM_STEPS_PER_FRAME', 10 if IN_DOCKER else 5))

# Overlay camera render resolution; keep auxiliary streams lean.
OVERLAY_RENDER_WIDTH = int(os.environ.get('OVERLAY_RENDER_WIDTH', 240))
OVERLAY_RENDER_HEIGHT = int(os.environ.get('OVERLAY_RENDER_HEIGHT', 160))

# Current robot type
current_robot = "g1"  # "g1", "spot", "ur5", or "ur5_t_push"
current_scene = None

# Environment instances (lazy loaded)
env_g1 = None
env_spot = None
env_ur5 = None
env_ur5_t_push = None
env = None  # Active environment

# Simulation state
running = True
latest_frame = None
frame_lock = threading.Lock()
mujoco_lock = threading.Lock()
renderer = None
needs_robot_switch = None  # Robot to switch to

episode_control_state = {
    "terminate": False,
    "truncate": False,
    "start_time": 0.0,  # Timestamp when episode started
    "max_duration": 120.0,  # Max episode duration in seconds (safety limit)
}

# Perception state - stores latest detected object poses
perception_state = {
    "poses": [],  # List of {object_id, position, orientation, confidence}
    "last_update": 0,  # Timestamp of last update
}
perception_lock = threading.Lock()
episode_control_lock = threading.Lock()

# Latest teleoperation action (for trainer state)
# Initialize with zero values instead of None so it's always a dict
last_teleop_action: dict[str, Any] = {
    "vx": 0.0,
    "vy": 0.0,
    "vz": 0.0,
    "vyaw": 0.0,
    "gripper": 0.0,  # 0-255 for UR5 gripper, 0 for robots without grippers
    # Cartesian rotation velocities (rad/s)
    "vrx": 0.0,
    "vry": 0.0,
    "vrz": 0.0,
    # Joint velocities (rad/s)
    "j1": 0.0,
    "j2": 0.0,
    "j3": 0.0,
    "j4": 0.0,
    "j5": 0.0,
    "j6": 0.0,
}
teleop_lock = threading.Lock()
last_teleop_command_time = 0.0  # Timestamp of last teleop command (for debugging)
jogging_active = False  # Track if jogging is currently active

# Homing state
homing_state = {
    "active": False,
    "current_joint": None,
    "joint_errors": [],
    "current_velocity": 0.5,  # Max velocity for homing
    "previous_error": None,
    "previous_abs_error": None,
    "stuck_count": 0,
    "is_jogging": False,
    "current_direction": None,
    "saved_velocity": None,  # Store velocity to restore after homing
}
homing_lock = threading.Lock()
jogging_pause_lock = threading.Lock()
jogging_paused = False

# WebSocket clients
ws_clients = set()
ws_clients_lock = threading.Lock()
external_ws_clients = set()  # External clients (not UI) that can send notifications
external_ws_clients_lock = threading.Lock()
client_metadata: dict = {}  # Metadata for all external clients
client_metadata_lock = threading.Lock()
# Camera state for orbit controls
cam = mujoco.MjvCamera()
cam.azimuth = 135
cam.elevation = -20
cam.distance = 3.0
cam.lookat = np.array([0, 0, 0.8])

# Camera follow mode
camera_follow = True

# Target visibility for push-T scene
t_target_visible = True

def _reset_camera_for_current_robot() -> None:
    """Reset the orbit camera to the defaults for the active robot."""
    cam.azimuth = 135
    cam.elevation = -20
    if current_robot == "g1":
        cam.distance = 3.0
        cam.lookat = np.array([0.0, 0.0, 0.8])
    elif current_robot in ("ur5", "ur5_t_push"):
        cam.distance = 1.8
        cam.lookat = np.array([0.3, 0.0, 0.6])
        cam.azimuth = 150
        cam.elevation = -25
    else:
        cam.distance = 2.5
        cam.lookat = np.array([0.0, 0.0, 0.4])


def _reset_active_environment() -> None:
    """Reset the primary simulation environment and associated camera."""
    with mujoco_lock:
        if env is not None:
            env.reset()
    _reset_camera_for_current_robot()
    # Track episode start time for duration watchdog
    with episode_control_lock:
        episode_control_state["start_time"] = time.time()
        episode_control_state["terminate"] = False
        episode_control_state["truncate"] = False


def _schedule_robot_switch(robot: str, scene: Optional[str]) -> None:
    """Schedule a robot/scene switch for the simulation loop."""
    global needs_robot_switch
    needs_robot_switch = {"robot": robot, "scene": scene}

MODEL_ROOT = Path(_nova_sim_dir) / "robots"

def discover_scene_files(robot: str) -> list[str]:
    """Return a list of available scene names for a robot."""
    model_dir = MODEL_ROOT / robot / "model"
    if not model_dir.exists():
        return []
    scenes = []
    for file in sorted(model_dir.glob("*.xml")):
        scenes.append(file.stem)
    return scenes

ROBOT_SCENES = {
    "ur5": ["scene_with_gripper"],
    "ur5_t_push": ["scene_t_push"],
    "g1": ["scene"],
    "spot": ["scene"],
}

ROBOT_LABELS = {
    "ur5": "Universal Robots UR5e",
    "ur5_t_push": "UR5e T-Push Scene",
    "g1": "Unitree G1 Humanoid",
    "spot": "Boston Dynamics Spot (Quadruped)",
}

DEFAULT_SCENES = {
    "ur5": "scene_with_gripper",
    "ur5_t_push": "scene_t_push",
    "g1": "scene",
    "spot": "scene",
}

AVAILABLE_ACTIONS = [
    "action",
    "reset",
    "switch_robot",
    "camera",
    "camera_follow",
    "toggle_target_visibility",
    "teleop_action",
    "arm_target",
    "gripper",
    "control_mode",
    "joint_positions",
    "arm_orientation",
    "use_orientation",
]

OVERLAY_CAMERA_PRESETS: dict[str, list[dict[str, Any]]] = {}

overlay_camera_states: dict[str, dict[str, Any]] = {}
overlay_frames: dict[str, bytes | None] = {}
overlay_frame_lock = threading.Lock()
overlay_camera_lock = threading.Lock()
dynamic_camera_configs: dict[str, dict[str, dict[str, Any]]] = {}

def _make_mjv_camera(config: dict[str, float]) -> mujoco.MjvCamera:
    cam_obj = mujoco.MjvCamera()
    cam_obj.lookat = np.array(config.get("lookat", [0.0, 0.0, 0.0]), dtype=np.float32)
    cam_obj.distance = float(config.get("distance", 1.0))
    cam_obj.azimuth = float(config.get("azimuth", 0.0))
    cam_obj.elevation = float(config.get("elevation", -30.0))
    return cam_obj


def _build_overlay_camera_state(config: dict[str, Any], create_renderer: bool = True) -> dict[str, Any]:
    """Build overlay camera state payload from config."""
    cam_obj = _make_mjv_camera(config)
    renderer_obj = None
    if create_renderer and env is not None:
        renderer_obj = mujoco.Renderer(
            env.model,
            height=OVERLAY_RENDER_HEIGHT,
            width=OVERLAY_RENDER_WIDTH
        )
    follow_site = config.get("follow_site")
    site_id = -1
    if follow_site and env is not None:
        site_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_SITE, follow_site)
    offset = np.array(config.get("offset", [0.0, 0.0, 0.0]), dtype=np.float32)
    state_payload = {
        "camera": cam_obj,
        "renderer": renderer_obj,
        "label": config.get("label", config["name"]),
        "follow_site": follow_site,
        "site_id": site_id,
        "offset": offset,
        "forward_offset": float(config.get("forward_offset", 0.0)),
        "track_orientation": bool(config.get("track_orientation")),
        "look_target_site": config.get("look_target_site"),
        "look_target_id": -1,
    }
    look_target_site = config.get("look_target_site")
    if look_target_site and env is not None:
        look_target_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_SITE, look_target_site)
        state_payload["look_target_id"] = look_target_id
    return state_payload


def _normalize_vec(vec: np.ndarray) -> np.ndarray:
    norm = np.linalg.norm(vec)
    if norm < 1e-8:
        return vec
    return vec / norm


def _rotation_matrix_to_quaternion(mat: np.ndarray) -> dict[str, float]:
    """Convert a 3x3 rotation matrix to quaternion (w, x, y, z)."""
    m00, m01, m02 = mat[0, 0], mat[0, 1], mat[0, 2]
    m10, m11, m12 = mat[1, 0], mat[1, 1], mat[1, 2]
    m20, m21, m22 = mat[2, 0], mat[2, 1], mat[2, 2]

    trace = m00 + m11 + m22
    if trace > 0.0:
        s = math.sqrt(trace + 1.0) * 2.0
        w = 0.25 * s
        x = (m21 - m12) / s
        y = (m02 - m20) / s
        z = (m10 - m01) / s
    elif m00 > m11 and m00 > m22:
        s = math.sqrt(1.0 + m00 - m11 - m22) * 2.0
        w = (m21 - m12) / s
        x = 0.25 * s
        y = (m01 + m10) / s
        z = (m02 + m20) / s
    elif m11 > m22:
        s = math.sqrt(1.0 + m11 - m00 - m22) * 2.0
        w = (m02 - m20) / s
        x = (m01 + m10) / s
        y = 0.25 * s
        z = (m12 + m21) / s
    else:
        s = math.sqrt(1.0 + m22 - m00 - m11) * 2.0
        w = (m10 - m01) / s
        x = (m02 + m20) / s
        y = (m12 + m21) / s
        z = 0.25 * s

    return {"w": float(w), "x": float(x), "y": float(y), "z": float(z)}


def _camera_axes_from_orbit(cam_obj: mujoco.MjvCamera) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return (right, up, forward) axes for an orbit camera in world coordinates."""
    azimuth = math.radians(float(cam_obj.azimuth))
    elevation = math.radians(float(cam_obj.elevation))

    right = np.array([math.sin(azimuth), -math.cos(azimuth), 0.0], dtype=np.float32)
    up = np.array([
        -math.cos(azimuth) * math.sin(elevation),
        -math.sin(azimuth) * math.sin(elevation),
        math.cos(elevation),
    ], dtype=np.float32)
    forward = np.array([
        -math.cos(azimuth) * math.cos(elevation),
        -math.sin(azimuth) * math.cos(elevation),
        -math.sin(elevation),
    ], dtype=np.float32)

    return _normalize_vec(right), _normalize_vec(up), _normalize_vec(forward)


def _camera_pose(cam_obj: mujoco.MjvCamera) -> dict[str, dict[str, float]]:
    """Compute camera pose (position + orientation) in world coordinates."""
    lookat = np.array(cam_obj.lookat, dtype=np.float32)
    right, up, forward = _camera_axes_from_orbit(cam_obj)
    position = lookat - forward * float(cam_obj.distance)
    rotation = np.column_stack((right, up, forward))
    quat = _rotation_matrix_to_quaternion(rotation)
    return {
        "position": {"x": float(position[0]), "y": float(position[1]), "z": float(position[2])},
        "orientation": quat,
    }


def _camera_fovy_degrees(cam_obj: mujoco.MjvCamera, env_obj) -> float:
    fovy = float(getattr(cam_obj, "fovy", 0.0) or 0.0)
    if fovy <= 0.0 and env_obj is not None:
        fovy = float(env_obj.model.vis.global_.fovy)
    return float(fovy) if fovy > 0.0 else 45.0


def _camera_intrinsics(width: int, height: int, fovy_degrees: float) -> dict[str, float]:
    """Compute pinhole intrinsics from image size and vertical field-of-view."""
    fovy_rad = math.radians(float(fovy_degrees))
    fy = (float(height) / 2.0) / math.tan(fovy_rad / 2.0)
    fx = fy
    return {
        "fx": float(fx),
        "fy": float(fy),
        "cx": float(width) / 2.0,
        "cy": float(height) / 2.0,
        "width": int(width),
        "height": int(height),
        "fovy_degrees": float(fovy_degrees),
    }


def _get_site_forward(env_obj, site_id: int) -> np.ndarray:
    """Compute the forward (X axis) vector for a site."""
    default = np.array([0.0, 0.0, 1.0], dtype=np.float32)
    if env_obj is None or site_id < 0:
        return default
    if site_id >= getattr(env_obj.model, "nsite", 0):
        return default
    xmat = getattr(env_obj.data, "site_xmat", None)
    if xmat is None or xmat.size < (site_id + 1) * 9:
        return default
    start = site_id * 9
    mat = xmat[start:start + 9]
    if mat.size < 9:
        return default
    try:
        forward = np.array([mat[0], mat[3], mat[6]], dtype=np.float32)
    except IndexError:
        return default
    norm = np.linalg.norm(forward)
    if norm < 1e-6:
        return default
    return forward / norm

def _close_overlay_renderers():
    global overlay_camera_states
    with overlay_camera_lock:
        for state in overlay_camera_states.values():
            renderer_obj = state.get("renderer")
            if renderer_obj:
                renderer_obj.close()
        overlay_camera_states = {}
    with overlay_frame_lock:
        overlay_frames.clear()


def _camera_scope_key(robot_type: str | None, scene_name: str | None) -> str:
    return f"{robot_type or ''}:{scene_name or ''}"


def _get_dynamic_camera_configs(robot_type: str | None, scene_name: str | None) -> list[dict[str, Any]]:
    key = _camera_scope_key(robot_type, scene_name)
    configs = dynamic_camera_configs.get(key, {})
    return list(configs.values())


def _store_dynamic_camera_config(robot_type: str | None, scene_name: str | None, config: dict[str, Any]) -> None:
    key = _camera_scope_key(robot_type, scene_name)
    if key not in dynamic_camera_configs:
        dynamic_camera_configs[key] = {}
    dynamic_camera_configs[key][config["name"]] = config


def prepare_overlay_renderers(robot_type: str, scene_name: str | None):
    """Create overlay renderers for the active robot/scene (if configured)."""
    _close_overlay_renderers()

    configs = []
    if robot_type and robot_type in OVERLAY_CAMERA_PRESETS:
        configs = OVERLAY_CAMERA_PRESETS[robot_type]
    elif scene_name and scene_name in OVERLAY_CAMERA_PRESETS:
        configs = OVERLAY_CAMERA_PRESETS[scene_name]
    dynamic_configs = _get_dynamic_camera_configs(robot_type, scene_name)
    if dynamic_configs:
        merged: dict[str, dict[str, Any]] = {}
        for config in configs:
            merged[config["name"]] = config
        for config in dynamic_configs:
            merged[config["name"]] = config
        configs = list(merged.values())
    if not configs:
        return

    for config in configs:
        state_payload = _build_overlay_camera_state(config)
        with overlay_camera_lock:
            overlay_camera_states[config["name"]] = state_payload
        with overlay_frame_lock:
            overlay_frames[config["name"]] = None


def _parse_bool(value):
    if value is None:
        return False
    return str(value).strip().lower() in ("1", "true", "yes", "on")


def _coerce_vec3(value: Any, default: list[float]) -> list[float]:
    if value is None:
        return default
    if isinstance(value, dict):
        value = [value.get("x"), value.get("y"), value.get("z")]
    if not isinstance(value, (list, tuple)) or len(value) != 3:
        raise ValueError("Expected a 3-element vector")
    return [float(value[0]), float(value[1]), float(value[2])]


def _valid_camera_name(name: str) -> bool:
    if not name:
        return False
    if name == "main":
        return False
    return re.match(r"^[A-Za-z0-9_-]+$", name) is not None


def _build_dynamic_camera_config(payload: dict[str, Any]) -> dict[str, Any]:
    name = str(payload.get("name", "")).strip()
    if not _valid_camera_name(name):
        raise ValueError("Invalid camera name. Use letters, numbers, hyphens, or underscores.")

    label = str(payload.get("label") or name).strip()
    default_lookat = [float(cam.lookat[0]), float(cam.lookat[1]), float(cam.lookat[2])]
    default_distance = float(cam.distance)
    default_azimuth = float(cam.azimuth)
    default_elevation = float(cam.elevation)

    lookat = _coerce_vec3(payload.get("lookat"), default_lookat)
    distance = float(payload.get("distance", default_distance))
    azimuth = float(payload.get("azimuth", default_azimuth))
    elevation = float(payload.get("elevation", default_elevation))

    config: dict[str, Any] = {
        "name": name,
        "label": label,
        "lookat": lookat,
        "distance": distance,
        "azimuth": azimuth,
        "elevation": elevation,
    }

    for key in ("follow_site", "look_target_site"):
        if payload.get(key) is not None:
            config[key] = str(payload.get(key))
    if payload.get("offset") is not None:
        config["offset"] = _coerce_vec3(payload.get("offset"), [0.0, 0.0, 0.0])
    if payload.get("forward_offset") is not None:
        config["forward_offset"] = float(payload.get("forward_offset", 0.0))
    if payload.get("track_orientation") is not None:
        config["track_orientation"] = bool(payload.get("track_orientation"))

    return config


def _load_nova_ur5_config_from_env():
    use_state = _parse_bool(os.environ.get("NOVA_UR5_USE_STATE_STREAM") or os.environ.get("NOVA_USE_STATE_STREAM"))
    use_ik = _parse_bool(os.environ.get("NOVA_UR5_USE_IK") or os.environ.get("NOVA_USE_IK"))
    if not (use_state or use_ik):
        return None

    instance_url = os.environ.get("NOVA_INSTANCE_URL")
    access_token = os.environ.get("NOVA_ACCESS_TOKEN")
    cell_id = os.environ.get("NOVA_CELL_ID")
    controller_id = os.environ.get("NOVA_CONTROLLER_ID")
    motion_group_id = os.environ.get("NOVA_MOTION_GROUP_ID")
    missing = [key for key, value in {
        "NOVA_INSTANCE_URL": instance_url,
        "NOVA_ACCESS_TOKEN": access_token,
        "NOVA_CELL_ID": cell_id,
        "NOVA_CONTROLLER_ID": controller_id,
        "NOVA_MOTION_GROUP_ID": motion_group_id,
    }.items() if not value]
    if missing:
        print(f"Nova API integration disabled; missing env vars: {', '.join(missing)}")
        return None

    return {
        "instance_url": instance_url,
        "access_token": access_token,
        "cell_id": cell_id,
        "controller_id": controller_id,
        "motion_group_id": motion_group_id,
        "motion_group_model": os.environ.get("NOVA_MOTION_GROUP_MODEL"),
        "tcp_name": os.environ.get("NOVA_TCP_NAME"),
        "response_rate_ms": int(os.environ.get("NOVA_RESPONSE_RATE_MS", "200")),
        "use_state_stream": use_state,
        "use_ik": use_ik,
    }


def init_g1():
    """Initialize G1 environment."""
    global env_g1
    if env_g1 is None:
        env_g1 = G1Env(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT)
        env_g1.scene_name = "scene"
        env_g1.reset()
        return env_g1
    env_g1.scene_name = "scene"
    return env_g1


def init_spot():
    """Initialize Spot environment with PyMPC controller."""
    global env_spot
    if env_spot is None:
        # Import SpotEnv from robots/spot directory
        spot_dir = os.path.join(_nova_sim_dir, 'robots', 'spot')
        sys.path.insert(0, spot_dir)
        from spot_env import SpotEnv
        sys.path.pop(0)
        # Use pympc controller if available, otherwise fallback to mpc_gait
        env_spot = SpotEnv(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT, controller_type='pympc')
        env_spot.scene_name = "scene"
        env_spot.reset()
        return env_spot
    env_spot.scene_name = "scene"
    return env_spot


def init_ur5(scene_name=DEFAULT_SCENES["ur5"]):
    """Initialize UR5e environment."""
    global env_ur5, env_ur5_t_push
    # Import UR5Env from robots/ur5 directory
    ur5_dir = os.path.join(_nova_sim_dir, 'robots', 'ur5')
    sys.path.insert(0, ur5_dir)
    from ur5_env import UR5Env
    sys.path.pop(0)

    # Check if Nova API is configured via environment variables
    nova_config = None
    if os.getenv('NOVA_INSTANCE_URL') and os.getenv('NOVA_ACCESS_TOKEN'):
        # Nova API is configured - attempt bidirectional control
        # State streaming: Mirror real robot pose while still running physics for objects (T-shape)
        # Jogging: UI can send directional jog commands to control the real robot
        # Note: If Nova connection fails, UR5Env will automatically fall back to internal IK
        print("Nova API credentials detected - attempting bidirectional control")
        nova_config = {
            "use_state_stream": True,  # Attempt to mirror real robot pose
            "use_ik": False,  # Use jogging API instead of Nova IK (falls back to internal IK if connection fails)
            "use_jogging": True  # Attempt bidirectional robot control
        }

    if scene_name == "scene_t_push":
        if env_ur5_t_push is None:
            env_ur5_t_push = UR5Env(
                render_mode="rgb_array",
                width=RENDER_WIDTH,
                height=RENDER_HEIGHT,
                scene_name="scene_t_push",
                nova_api_config=nova_config
            )
            env_ur5_t_push.reset()
        return env_ur5_t_push

    if env_ur5 is None:
        env_ur5 = UR5Env(
            render_mode="rgb_array",
            width=RENDER_WIDTH,
            height=RENDER_HEIGHT,
            nova_api_config=nova_config
        )
        env_ur5.reset()
    return env_ur5


def switch_robot(robot_type, scene_name=None):
    """Switch to a different robot."""
    global env, current_robot, current_scene, renderer, cam

    active_scene = None
    if robot_type == current_robot and env is not None:
        return


    # Close old renderer
    if renderer is not None:
        renderer.close()
        renderer = None

    if robot_type == "g1":
        env = init_g1()
        cam.lookat = np.array([0, 0, 0.8])
        cam.distance = 3.0
    elif robot_type == "spot":
        env = init_spot()
        cam.lookat = np.array([0, 0, 0.4])
        cam.distance = 2.5
    elif robot_type == "ur5":
        selected_scene = scene_name or DEFAULT_SCENES.get("ur5")
        env = init_ur5(selected_scene)
        active_scene = selected_scene
        cam.lookat = np.array([0.3, 0, 0.6])
        cam.distance = 1.8
        cam.azimuth = 150
        cam.elevation = -25
    elif robot_type == "ur5_t_push":
        selected_scene = scene_name or "scene_t_push"
        env = init_ur5(selected_scene)
        active_scene = selected_scene
        cam.lookat = np.array([0.5, 0, 0.55])
        cam.distance = 1.9
        cam.azimuth = 150
        cam.elevation = -25
    else:
        print(f"Unknown robot type: {robot_type}")
        return

    current_robot = robot_type
    current_scene = active_scene
    env.reset()

    # Track episode start time for duration watchdog
    with episode_control_lock:
        episode_control_state["start_time"] = time.time()
        episode_control_state["terminate"] = False
        episode_control_state["truncate"] = False

    # Initialize gripper value in teleop_action for UR5
    with teleop_lock:
        if robot_type in ("ur5", "ur5_t_push"):
            has_gripper = getattr(env, "has_gripper", False)
            if has_gripper:
                # Get current gripper value from env
                gripper_val = getattr(env, "get_gripper", lambda: 128)()
                last_teleop_action["gripper"] = float(gripper_val)
            else:
                last_teleop_action["gripper"] = 0.0
        else:
            # Locomotion robots don't have grippers
            last_teleop_action["gripper"] = 0.0

    # Create new renderer
    renderer = mujoco.Renderer(env.model, height=env.height, width=env.width)

    prepare_overlay_renderers(robot_type, active_scene)


# Initialize with G1
env = init_g1()


def _get_scene_objects():
    """Extract scene objects (position and orientation) from the environment."""
    if env is None:
        return []

    scene_objects = []

    # For UR5 scenes, extract T-shape objects if they exist
    if current_robot in ("ur5", "ur5_t_push"):
        # Get t_object body if it exists
        t_object_id = getattr(env, 't_object_body_id', -1)
        if t_object_id >= 0:
            pos = env.data.xpos[t_object_id]
            # MuJoCo stores quaternions in data.xquat as [w, x, y, z]
            quat = env.data.xquat[t_object_id]
            scene_objects.append({
                'name': 't_object',
                'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
                'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
            })

        # Get t_target body if it exists
        t_target_id = getattr(env, 't_target_body_id', -1)
        if t_target_id >= 0:
            pos = env.data.xpos[t_target_id]
            quat = env.data.xquat[t_target_id]
            scene_objects.append({
                'name': 't_target',
                'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
                'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
            })

        # Get box body if it exists
        try:
            box_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, "box")
            if box_id >= 0:
                pos = env.data.xpos[box_id]
                quat = env.data.xquat[box_id]
                scene_objects.append({
                    'name': 'box',
                    'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
                    'orientation': {'w': float(quat[0]), 'x': float(quat[1]), 'y': float(quat[2]), 'z': float(quat[3])}
                })
        except:
            pass

    return scene_objects


def broadcast_state():
    """Broadcast robot state to all connected WebSocket clients."""
    with mujoco_lock:
        if env is None:
            return
        obs = env._get_obs()
        cmd = env.get_command()
        steps = env.steps
        command = {
            "vx": float(cmd[0]) if len(cmd) > 0 else 0.0,
            "vy": float(cmd[1]) if len(cmd) > 1 else 0.0,
            "vyaw": float(cmd[2]) if len(cmd) > 2 else 0.0,
            "vz": 0.0,
        }
        with teleop_lock:
            teleop_snapshot = last_teleop_action.copy()
        reward_value = None
        if hasattr(env, "get_task_reward"):
            reward_value = env.get_task_reward()

        # Build lists of identified and unidentified clients
        with client_metadata_lock:
            identified_clients = []
            unidentified_clients = []
            for ws, info in list(client_metadata.items()):
                identity = info.get("identity")
                if identity:
                    # Client has identified themselves
                    if ws in external_ws_clients:
                        identified_clients.append(identity)
                else:
                    # Client connected but hasn't sent identity yet
                    unidentified_clients.append(f"unidentified-{id(ws) % 10000}")

        # Get scene objects
        scene_objects = _get_scene_objects()

        # UR5 has different state structure
        if current_robot in ("ur5", "ur5_t_push"):
            ee_pos = env.get_end_effector_pos()
            ee_quat = env.get_end_effector_quat()
            target = env.get_target()
            target_euler = env.get_target_orientation()
            gripper = env.get_gripper()
            joint_pos = env.get_joint_positions()
            joint_targets = env.get_joint_targets()
            control_mode = env.get_control_mode()

            # Check Nova API status
            nova_client = getattr(env, '_nova_client', None)
            nova_state_streaming = getattr(env, '_use_nova_state_stream', False)
            nova_ik = getattr(env, '_use_nova_ik', False)
            # Check if actually connected (WebSocket active and receiving data)
            nova_connected = False
            if nova_client is not None:
                if hasattr(nova_client, 'is_state_stream_connected'):
                    # Check actual stream connection status
                    nova_connected = nova_client.is_state_stream_connected()
                elif hasattr(nova_client, 'is_connected'):
                    # Fallback: check general connection status
                    nova_connected = nova_client.is_connected()
                else:
                    # If client exists but we can't verify connection, assume connected
                    # This handles the IK-only case where state streaming isn't used
                    nova_connected = True

            nova_available = False
            nova_enabled_pref = False
            has_nova_conf = getattr(env, 'has_nova_configuration', None)
            if callable(has_nova_conf):
                nova_available = has_nova_conf()
            is_nova_enabled = getattr(env, 'is_nova_enabled', None)
            if callable(is_nova_enabled):
                nova_enabled_pref = is_nova_enabled()


            state_msg = json.dumps({
                'type': 'state',
                'data': {
                    'robot': current_robot,
                    'scene': current_scene,
                    'observation': {
                        'end_effector': {'x': float(ee_pos[0]), 'y': float(ee_pos[1]), 'z': float(ee_pos[2])},
                        'ee_orientation': {'w': float(ee_quat[0]), 'x': float(ee_quat[1]), 'y': float(ee_quat[2]), 'z': float(ee_quat[3])},
                        'ee_target': {'x': float(target[0]), 'y': float(target[1]), 'z': float(target[2])},
                        'ee_target_orientation': {'roll': float(target_euler[0]), 'pitch': float(target_euler[1]), 'yaw': float(target_euler[2])},
                        'gripper': float(gripper),
                        'joint_positions': [float(j) for j in joint_pos],
                        'joint_targets': [float(j) for j in joint_targets],
                    },
                    'control_mode': control_mode,
                    'steps': int(steps),
                    'reward': reward_value,
                    'teleop_action': teleop_snapshot,
                    'scene_objects': scene_objects,
                    'nova_api': {
                        'connected': nova_connected,
                        'state_streaming': nova_state_streaming,
                        'ik': nova_ik,
                        'available': nova_available,
                        'enabled': nova_enabled_pref
                    },
                    'connected_clients': identified_clients,
                    'unidentified_clients': unidentified_clients
                }
            })
        else:
            # Locomotion robots (G1, Spot)
            base_pos = obs[0:3]  # x, y, z position
            base_quat = obs[3:7]  # quaternion orientation (w, x, y, z)

            state_msg = json.dumps({
                'type': 'state',
                'data': {
                    'robot': current_robot,
                    'scene': current_scene,
                    'observation': {
                        'position': {'x': float(base_pos[0]), 'y': float(base_pos[1]), 'z': float(base_pos[2])},
                        'orientation': {'w': float(base_quat[0]), 'x': float(base_quat[1]), 'y': float(base_quat[2]), 'z': float(base_quat[3])},
                    },
                    'steps': int(steps),
                    'reward': reward_value,
                    'teleop_action': teleop_snapshot,
                    'scene_objects': scene_objects,
                    'connected_clients': identified_clients,
                    'unidentified_clients': unidentified_clients
                }
            })

    # Send to all connected clients
        with ws_clients_lock:
            dead_clients = set()
            for ws in ws_clients:
                try:
                    ws.send(state_msg)
                except Exception:
                    dead_clients.add(ws)
        # Remove dead clients
        ws_clients.difference_update(dead_clients)


def broadcast_camera_event(action: str, camera_payload: dict[str, Any], scope: dict[str, Any]) -> None:
    """Announce camera changes via the state stream."""
    message = {
        "type": "state",
        "data": {
            "camera_event": {
                "action": action,
                "camera": camera_payload,
                "scope": scope,
                "timestamp": time.time(),
            }
        }
    }
    with ws_clients_lock:
        dead_clients = set()
        for ws in ws_clients:
            try:
                ws.send(json.dumps(message))
            except Exception:
                dead_clients.add(ws)
        ws_clients.difference_update(dead_clients)


def _handle_external_client_message(ws, data):
    """Process message payloads originating from external clients."""
    msg_type = data.get("type")
    if msg_type in ("client_identity", "trainer_identity"):  # Support old name for backward compat
        payload = data.get("data", {}) or {}
        identity = payload.get("client_id") or payload.get("trainer_id") or payload.get("name") or "client"
        with external_ws_clients_lock:
            external_ws_clients.add(ws)
        # Don't call _register_external_client again - client is already registered on connection
        _set_client_identity(ws, identity)
        broadcast_state()
        broadcast_connected_clients_status()
        return
    if msg_type in ("client_notification", "notification"):  # Support old name for backward compat
        payload = data.get("data", {})
        payload.setdefault("timestamp", time.time())
        broadcast_notification_to_all(payload)


def _build_connected_clients_payload():
    """Build a summary payload describing connected external clients."""
    with client_metadata_lock:
        identified = []
        unidentified = []
        for ws, info in client_metadata.items():
            identity = info.get("identity")
            if identity:
                # Client has identified themselves
                if ws in external_ws_clients:
                    identified.append(identity)
            else:
                # Client connected but hasn't sent identity yet
                unidentified.append(f"unidentified-{id(ws) % 10000}")

    payload = {
        "clients": identified,
        "unidentified_clients": unidentified,
    }
    return payload


def _register_external_client(ws):
    with client_metadata_lock:
        client_metadata[ws] = {"identity": None, "connected_at": time.time()}


def _unregister_external_client(ws):
    with client_metadata_lock:
        client_metadata.pop(ws, None)


def _set_client_identity(ws, identity: str):
    with client_metadata_lock:
        info = client_metadata.get(ws)
        if not info:
            # Client wasn't registered on connection (shouldn't happen with new flow)
            # Register them now for backward compatibility
            client_metadata[ws] = {"identity": identity, "connected_at": time.time(), "last_seen": time.time()}
            return
        info["identity"] = identity
        info["last_seen"] = time.time()


def broadcast_connected_clients_status():
    """Notify all clients about connected external clients."""
    payload = _build_connected_clients_payload()
    message = json.dumps({"type": "connected_clients", "data": payload})
    with ws_clients_lock:
        dead_clients = set()
        for client in ws_clients:
            try:
                client.send(message)
            except:
                dead_clients.add(client)
        ws_clients.difference_update(dead_clients)

def broadcast_to_external_clients(event_type: str, payload: dict):
    """Broadcast a structured event to connected external WebSocket clients."""
    message = json.dumps({"type": event_type, "data": payload})
    with external_ws_clients_lock:
        dead_clients = set()
        for ws in external_ws_clients:
            try:
                ws.send(message)
            except Exception:
                dead_clients.add(ws)
        external_ws_clients.difference_update(dead_clients)


def broadcast_notification_to_all(payload: dict):
    """Send a notification message (from external client) to all clients."""
    notification = json.dumps({"type": "client_notification", "data": payload})
    with ws_clients_lock:
        dead_clients = set()
        for ws in ws_clients:
            try:
                ws.send(notification)
            except Exception:
                dead_clients.add(ws)
        ws_clients.difference_update(dead_clients)


def _safe_ws_send(ws, message: dict):
    """Send JSON message over WebSocket without raising."""
    try:
        ws.send(json.dumps(message))
    except Exception:
        pass



def _signal_episode_control(action: str):
    """Set episode control flags and notify trainer/UI clients."""
    action = (action or "").lower()
    if action not in ("terminate", "truncate"):
        return

    with episode_control_lock:
        episode_control_state[action] = True
        episode_control_state["start_time"] = 0.0  # Reset timer
    status_label = "terminated" if action == "terminate" else "truncated"
    payload = {
        "action": action,
        "status": status_label,
        "source": "ui",
        "timestamp": time.time(),
    }
    broadcast_to_external_clients("episode_event", payload)
    # Reset UI environment so the user can continue without manual reset
    with mujoco_lock:
        if env is not None:
            env.reset()
    # Track new episode start time
    with episode_control_lock:
        episode_control_state["start_time"] = time.time()
        episode_control_state["terminate"] = False
        episode_control_state["truncate"] = False


def _consume_episode_control_flags():
    """Return and clear the active episode control flags."""
    with episode_control_lock:
        terminate = episode_control_state.get("terminate", False)
        truncate = episode_control_state.get("truncate", False)
        episode_control_state["terminate"] = False
        episode_control_state["truncate"] = False
    return terminate, truncate


def _quat_to_mat(quat):
    """Convert quaternion (w, x, y, z) to 3x3 rotation matrix.
    
    Args:
        quat: dict with keys 'w', 'x', 'y', 'z' or tuple/list [w, x, y, z]
    
    Returns:
        3x3 rotation matrix as numpy array
    """
    if isinstance(quat, dict):
        w, x, y, z = quat['w'], quat['x'], quat['y'], quat['z']
    else:
        w, x, y, z = quat
    
    # Normalize quaternion
    norm = math.sqrt(w*w + x*x + y*y + z*z)
    if norm < 1e-10:
        return np.eye(3)
    w, x, y, z = w/norm, x/norm, y/norm, z/norm
    
    # Convert to rotation matrix
    mat = np.array([
        [1 - 2*(y*y + z*z),     2*(x*y - w*z),     2*(x*z + w*y)],
        [    2*(x*y + w*z), 1 - 2*(x*x + z*z),     2*(y*z - w*x)],
        [    2*(x*z - w*y),     2*(y*z + w*x), 1 - 2*(x*x + y*y)]
    ])
    return mat


def _render_perception_frames(renderer_obj, model, data):
    """Render bounding boxes for detected object poses.

    Draws wireframe boxes at detected poses, color-coded by confidence.
    Green (high confidence >0.8), Yellow (medium 0.5-0.8), Red (low <0.5).
    """
    with perception_lock:
        poses = perception_state.get("poses", [])
        last_update = perception_state.get("last_update", 0)

    # Don't render if data is too old (> 1 second)
    if time.time() - last_update > 1.0:
        return

    # Render bounding box for each detected pose
    for pose in poses:
        try:
            pos = pose.get("position", {})
            orientation = pose.get("orientation", {})
            confidence = pose.get("confidence", 0.0)
            dimensions = pose.get("dimensions")

            # Skip if no dimensions provided
            if not dimensions:
                continue

            # Extract position
            px = float(pos.get("x", 0.0))
            py = float(pos.get("y", 0.0))
            pz = float(pos.get("z", 0.0))

            # Extract dimensions
            if isinstance(dimensions, dict):
                width = float(dimensions.get("width", 0.05))
                height = float(dimensions.get("height", 0.05))
                depth = float(dimensions.get("depth", 0.05))
            elif isinstance(dimensions, (list, tuple)):
                width, height, depth = float(dimensions[0]), float(dimensions[1]), float(dimensions[2])
            else:
                continue

            # Color based on confidence: high=green, medium=yellow, low=red
            if confidence > 0.8:
                rgba = [0.0, 1.0, 0.0, 0.5]  # Green
            elif confidence > 0.5:
                rgba = [1.0, 1.0, 0.0, 0.5]  # Yellow
            else:
                rgba = [1.0, 0.0, 0.0, 0.5]  # Red

            # Convert quaternion to rotation matrix
            rot_mat = _quat_to_mat(orientation)

            # Render as box geometry
            # MuJoCo box size is half-extents, so divide by 2
            mujoco.mjv_initGeom(
                renderer_obj.scene.geoms[renderer_obj.scene.ngeom],
                type=mujoco.mjtGeom.mjGEOM_BOX,
                size=[width/2, height/2, depth/2],
                pos=[px, py, pz],
                mat=rot_mat.flatten(),
                rgba=rgba
            )
            renderer_obj.scene.ngeom += 1

            # Limit to max geoms to avoid overflow
            if renderer_obj.scene.ngeom >= renderer_obj.scene.maxgeom - 10:
                break

        except Exception as e:
            # Silently skip invalid poses
            pass




def simulation_loop():
    global running, latest_frame, camera_follow, t_target_visible, renderer, needs_robot_switch, env

    print("Starting simulation loop...")

    # Initialize renderer
    renderer = mujoco.Renderer(env.model, height=env.height, width=env.width)

    physics_steps_per_frame = 5
    broadcast_counter = 0

    def _run_iteration():
        nonlocal broadcast_counter
        global needs_robot_switch, latest_frame, jogging_active, last_teleop_command_time

        if needs_robot_switch is not None:
            with mujoco_lock:
                switch_robot(
                    needs_robot_switch.get("robot"),
                    scene_name=needs_robot_switch.get("scene")
                )
            needs_robot_switch = None

        # Check episode duration watchdog
        current_time = time.time()
        with episode_control_lock:
            episode_start = episode_control_state.get("start_time", 0.0)
            max_duration = episode_control_state.get("max_duration", 120.0)
            if episode_start > 0 and (current_time - episode_start) > max_duration:
                # Episode has exceeded max duration - automatically truncate
                print(f"[Watchdog] Episode exceeded max duration ({max_duration}s) - truncating episode")
                episode_control_state["truncate"] = True
                episode_control_state["start_time"] = 0.0  # Reset timer
                # Stop any active jogging
                if jogging_active:
                    with mujoco_lock:
                        if env is not None and current_robot in ("ur5", "ur5_t_push"):
                            stop_jog_fn = getattr(env, "stop_jog", None)
                            if callable(stop_jog_fn):
                                try:
                                    stop_jog_fn()
                                    jogging_active = False
                                except Exception as e:
                                    print(f"[Watchdog] Error stopping jogging during truncation: {e}")

        with mujoco_lock:
            if env is None:
                return

            sim_dt = env.model.opt.timestep

            # Run physics steps
            for _ in range(physics_steps_per_frame):
                env.step_with_controller(dt=sim_dt)

            # Update camera to follow robot or object
            if current_robot in ("ur5", "ur5_t_push"):
                # For UR5: follow t_object when enabled, table center when disabled
                if camera_follow:
                    # Follow t_object if it exists
                    t_object_id = getattr(env, 't_object_body_id', -1)
                    if t_object_id >= 0:
                        object_pos = env.data.xpos[t_object_id]
                        cam.lookat[0] = object_pos[0]
                        cam.lookat[1] = object_pos[1]
                        cam.lookat[2] = object_pos[2]
                else:
                    # Center on table
                    cam.lookat[0] = 0.5
                    cam.lookat[1] = 0.0
                    cam.lookat[2] = 0.42
            elif camera_follow:
                # For locomotion robots: follow robot base
                robot_pos = env.data.qpos[:3]
                cam.lookat[0] = robot_pos[0]
                cam.lookat[1] = robot_pos[1]
                cam.lookat[2] = 0.8 if current_robot == "g1" else 0.4

            # Render
            if renderer is not None:
                renderer.update_scene(env.data, camera=cam)

                # Hide t_target if visibility is disabled (push-T scene only)
                if current_robot == "ur5_t_push" and not t_target_visible:
                    t_target_id = getattr(env, 't_target_body_id', -1)
                    if t_target_id >= 0:
                        # Set alpha to 0 for all geoms belonging to t_target body
                        for i in range(renderer.scene.ngeom):
                            geom = renderer.scene.geoms[i]
                            # Check if this geom belongs to t_target body by comparing position
                            target_pos = env.data.xpos[t_target_id]
                            geom_pos = geom.pos
                            # If geom is very close to target body, it likely belongs to it
                            dist = np.linalg.norm(np.array(geom_pos) - target_pos)
                            if dist < 0.15:  # Within 15cm - likely part of the T-shape
                                geom.rgba[3] = 0.0  # Set alpha to 0 (invisible)

                # Add perception visualization markers
                _render_perception_frames(renderer, env.model, env.data)
                frame = renderer.render()
            else:
                return

        # Encode frame
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        ret, buffer = cv2.imencode('.jpg', frame_bgr)

        if ret:
            with frame_lock:
                latest_frame = buffer.tobytes()

        # Render overlay camera frames (if configured)
        with overlay_camera_lock:
            overlay_snapshot = list(overlay_camera_states.items())
        for name, state in overlay_snapshot:
            renderer_obj = state.get("renderer")
            cam_obj = state.get("camera")
            if renderer_obj is None or cam_obj is None:
                if cam_obj is None:
                    continue
                with mujoco_lock:
                    if env is None:
                        continue
                    try:
                        renderer_obj = mujoco.Renderer(
                            env.model,
                            height=OVERLAY_RENDER_HEIGHT,
                            width=OVERLAY_RENDER_WIDTH
                        )
                    except Exception:
                        continue
                with overlay_camera_lock:
                    if name in overlay_camera_states:
                        overlay_camera_states[name]["renderer"] = renderer_obj
                if renderer_obj is None:
                    continue
            if state.get("follow_site") and env is not None:
                site_id = state.get("site_id", -1)
                if site_id >= 0 and site_id < getattr(env.model, "nsite", 0):
                    site_xpos = getattr(env.data, "site_xpos", None)
                    if site_xpos is not None and site_xpos.size >= (site_id + 1) * 3:
                        lookat_point = site_xpos[site_id]
                    else:
                        continue
                    offset = state.get("offset")
                    if offset is not None:
                        lookat_point = lookat_point + offset
                    forward_offset = state.get("forward_offset", 0.0)
                    if forward_offset:
                        forward = _get_site_forward(env, site_id)
                        lookat_point = lookat_point + forward * forward_offset
                    look_target_id = state.get("look_target_id", -1)
                    target_point = None
                    if look_target_id >= 0 and site_xpos.size >= (look_target_id + 1) * 3:
                        target_point = site_xpos[look_target_id]
                    cam_obj.lookat = target_point if target_point is not None else lookat_point
                    if state.get("track_orientation"):
                        forward = _get_site_forward(env, site_id)
                        yaw = math.degrees(math.atan2(forward[1], forward[0]))
                        pitch = math.degrees(math.asin(np.clip(forward[2], -1.0, 1.0)))
                        cam_obj.azimuth = -yaw + 90.0
                        cam_obj.elevation = -pitch
            renderer_obj.update_scene(env.data, camera=cam_obj)

            # Hide t_target if visibility is disabled (push-T scene only)
            if current_robot == "ur5_t_push" and not t_target_visible:
                t_target_id = getattr(env, 't_target_body_id', -1)
                if t_target_id >= 0:
                    # Set alpha to 0 for all geoms belonging to t_target body
                    for i in range(renderer_obj.scene.ngeom):
                        geom = renderer_obj.scene.geoms[i]
                        # Check if this geom belongs to t_target body by comparing position
                        target_pos = env.data.xpos[t_target_id]
                        geom_pos = geom.pos
                        # If geom is very close to target body, it likely belongs to it
                        dist = np.linalg.norm(np.array(geom_pos) - target_pos)
                        if dist < 0.15:  # Within 15cm - likely part of the T-shape
                            geom.rgba[3] = 0.0  # Set alpha to 0 (invisible)

            overlay_frame = renderer_obj.render()
            if overlay_frame is None:
                continue
            overlay_bgr = cv2.cvtColor(overlay_frame, cv2.COLOR_RGB2BGR)
            ret2, buffer2 = cv2.imencode('.jpg', overlay_bgr)
            if ret2:
                with overlay_frame_lock:
                    overlay_frames[name] = buffer2.tobytes()

        # Broadcast state
        broadcast_counter += 1
        if broadcast_counter >= 6:
            broadcast_counter = 0
            broadcast_state()

    while running:
        loop_start = time.time()

        try:
            _run_iteration()
        except Exception as e:
            print(f"Error in simulation loop: {e}")
            import traceback
            traceback.print_exc()

        elapsed = time.time() - loop_start
        wait = max(0.001, 0.016 - elapsed)
        time.sleep(wait)

    if renderer:
        renderer.close()
    print("Simulation loop stopped.")


# Start simulation thread
sim_thread = threading.Thread(target=simulation_loop, daemon=True)
sim_thread.start()


def generate_frames():
    while True:
        with frame_lock:
            if latest_frame is not None:
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + latest_frame + b'\r\n')
        time.sleep(0.04)


def generate_overlay_frames(name: str):
    while True:
        with overlay_frame_lock:
            frame = overlay_frames.get(name)
        if frame is not None:
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
        time.sleep(0.04)


def _set_jogging_paused(paused: bool) -> bool:
    """Pause or resume jogging commands; returns the previous pause state."""
    global jogging_paused
    with jogging_pause_lock:
        previous = jogging_paused
        jogging_paused = paused
    return previous


def _is_jogging_paused() -> bool:
    with jogging_pause_lock:
        return jogging_paused


def _step_multi_axis_homing_locked(tolerance: float = 0.01) -> dict:
    """Run one homing control step (expects mujoco_lock + homing_lock held)."""
    homing_state["last_home_time"] = time.time()

    if env is None or current_robot not in ("ur5", "ur5_t_push"):
        return {"status": "unsupported"}

    home_pose = getattr(env, "home_pose", None)
    get_joints = getattr(env, "get_joint_positions", None)
    start_multi_joint_jog = getattr(env, "start_multi_joint_jog", None)
    stop_jog = getattr(env, "stop_jog", None)

    if home_pose is None or not callable(get_joints) or not callable(start_multi_joint_jog) or not callable(stop_jog):
        return {"status": "unsupported"}

    if hasattr(home_pose, "tolist"):
        target_positions = home_pose.tolist()
    else:
        target_positions = list(home_pose)

    current_positions = get_joints()

    # Calculate errors for all joints
    errors = []
    for joint_idx in range(min(len(target_positions), len(current_positions))):
        current = current_positions[joint_idx]
        target = target_positions[joint_idx]
        error = target - current
        errors.append(error)

    max_abs_error = max(abs(e) for e in errors) if errors else 0.0
    if max_abs_error < tolerance:
        if homing_state.get("is_jogging", False):
            stop_jog()
            homing_state["is_jogging"] = False
            homing_state["active"] = False
            homing_state.pop("previous_errors", None)
            homing_state.pop("current_velocities", None)
            print("[Homing] Completed - all joints at target")
        return {
            "status": "done",
            "errors": errors,
            "max_abs_error": max_abs_error,
        }

    if "previous_errors" not in homing_state:
        homing_state["previous_errors"] = [0.0] * len(errors)

    base_velocity = 0.15  # rad/s - slower to prevent overshoot
    velocities = []
    stop_tolerance = tolerance * 0.5

    for joint_idx, error in enumerate(errors):
        abs_error = abs(error)

        prev_error = homing_state["previous_errors"][joint_idx]
        if prev_error != 0 and ((prev_error > 0 and error < 0) or (prev_error < 0 and error > 0)):
            velocities.append(0.0)
            continue

        if abs_error < stop_tolerance:
            velocities.append(0.0)
        else:
            if max_abs_error > 0:
                velocity = base_velocity * (abs_error / max_abs_error)
            else:
                velocity = 0.0

            damping_zone = tolerance * 3.0
            if abs_error < damping_zone:
                damping_factor = (abs_error / damping_zone) ** 2
                velocity *= damping_factor

            min_velocity = 0.02
            if velocity < min_velocity:
                velocities.append(0.0)
            else:
                velocities.append(velocity if error > 0 else -velocity)

    homing_state["previous_errors"] = errors

    if all(abs(v) < 0.001 for v in velocities):
        if homing_state.get("is_jogging", False):
            stop_jog()
            homing_state["is_jogging"] = False
            homing_state["active"] = False
            homing_state.pop("previous_errors", None)
            homing_state.pop("current_velocities", None)
            print("[Homing] All joints reached target or stopped")
        return {
            "status": "done",
            "errors": errors,
            "max_abs_error": max_abs_error,
        }

    current_velocities = homing_state.get("current_velocities")
    velocities_changed = current_velocities is None or any(
        abs(v1 - v2) > 0.01 for v1, v2 in zip(velocities, current_velocities)
    )

    if velocities_changed:
        if homing_state.get("is_jogging", False):
            stop_jog()

        start_multi_joint_jog(velocities)
        homing_state["is_jogging"] = True
        homing_state["current_velocities"] = velocities
        homing_state["active"] = True

    return {
        "status": "running",
        "errors": errors,
        "max_abs_error": max_abs_error,
        "velocities": velocities,
    }


def handle_ws_message(ws, data):
    """Handle incoming WebSocket message."""
    global needs_robot_switch, camera_follow, t_target_visible, last_teleop_action, last_teleop_command_time, jogging_active

    msg_type = data.get('type')

    if msg_type in ("client_identity", "trainer_identity", "client_notification", "notification"):
        _handle_external_client_message(ws, data)
        return

    # Backward compatibility: map old message types to new ones
    if msg_type == 'command':
        msg_type = 'action'
        data['type'] = 'action'
    elif msg_type == 'teleop_command':
        msg_type = 'teleop_action'
        data['type'] = 'teleop_action'

    if msg_type == 'action':
        payload = data.get('data', {})
        vx = payload.get('vx', 0.0)
        vy = payload.get('vy', 0.0)
        vz = payload.get('vz', 0.0)
        vyaw = payload.get('vyaw', 0.0)
        vrx = payload.get('vrx', 0.0)
        vry = payload.get('vry', 0.0)
        vrz = payload.get('vrz', 0.0)
        j1 = payload.get('j1', 0.0)
        j2 = payload.get('j2', 0.0)
        j3 = payload.get('j3', 0.0)
        j4 = payload.get('j4', 0.0)
        j5 = payload.get('j5', 0.0)
        j6 = payload.get('j6', 0.0)
        gripper = payload.get('gripper', None)

        if current_robot in ("ur5", "ur5_t_push") and _is_jogging_paused():
            if gripper is not None:
                with mujoco_lock:
                    if env is not None:
                        env.set_gripper(float(gripper))
            return

        with mujoco_lock:
            if env is not None:
                # For UR5: translate velocity actions to jogging commands
                if current_robot in ("ur5", "ur5_t_push"):
                    # Handle gripper if specified
                    if gripper is not None:
                        env.set_gripper(float(gripper))

                    # Check which type of velocity is active (joint vs cartesian)
                    joint_velocities = [float(j1), float(j2), float(j3), float(j4), float(j5), float(j6)]
                    cartesian_translation = [float(vx), float(vy), float(vz)]
                    cartesian_rotation = [float(vrx), float(vry), float(vrz)]

                    any_joint = any(abs(v) > 0.001 for v in joint_velocities)
                    any_cartesian = any(abs(v) > 0.001 for v in cartesian_translation + cartesian_rotation)

                    # Update teleop command timestamp
                    last_teleop_command_time = time.time()

                    # Apply the appropriate jogging mode
                    if any_joint:
                        env.start_multi_joint_jog(joint_velocities)
                        jogging_active = True
                    elif any_cartesian:
                        # Convert m/s to mm/s for translations
                        translation_mm_s = [v * 1000.0 for v in cartesian_translation]
                        env.start_jog(
                            'cartesian_velocity',
                            translation=translation_mm_s,
                            rotation=cartesian_rotation,
                            tcp_id='Flange',
                            coord_system_id='world'
                        )
                        jogging_active = True
                    else:
                        # No active velocity - stop jogging
                        env.stop_jog()
                        jogging_active = False
                else:
                    # For locomotion robots: use vx, vy, vyaw
                    env.set_command(vx, vy, vyaw)

        with teleop_lock:
            # Update all velocity fields in teleop_action
            last_teleop_action["vx"] = float(vx)
            last_teleop_action["vy"] = float(vy)
            last_teleop_action["vz"] = float(vz)
            last_teleop_action["vyaw"] = float(vyaw)
            last_teleop_action["vrx"] = float(vrx)
            last_teleop_action["vry"] = float(vry)
            last_teleop_action["vrz"] = float(vrz)
            last_teleop_action["j1"] = float(j1)
            last_teleop_action["j2"] = float(j2)
            last_teleop_action["j3"] = float(j3)
            last_teleop_action["j4"] = float(j4)
            last_teleop_action["j5"] = float(j5)
            last_teleop_action["j6"] = float(j6)
            if gripper is not None:
                last_teleop_action["gripper"] = float(gripper)

    elif msg_type == 'reset':
        payload = data.get('data', {})
        seed = payload.get('seed')
        with mujoco_lock:
            if env is not None:
                obs, info = env.reset(seed=seed)
        # Track episode start time for duration watchdog
        with episode_control_lock:
            episode_control_state["start_time"] = time.time()
            episode_control_state["terminate"] = False
            episode_control_state["truncate"] = False

    elif msg_type == 'switch_robot':
        payload = data.get('data', {})
        robot = payload.get('robot', 'g1')
        scene = payload.get('scene')
        _schedule_robot_switch(robot, scene)

    elif msg_type == 'camera':
        payload = data.get('data', {})
        a = cam.azimuth * np.pi / 180.0
        e = cam.elevation * np.pi / 180.0

        cam_type = payload.get('action')
        if cam_type == 'rotate':
            cam.azimuth -= payload.get('dx', 0) * 0.5
            cam.elevation -= payload.get('dy', 0) * 0.5
            cam.elevation = np.clip(cam.elevation, -89, 89)
        elif cam_type == 'zoom':
            cam.distance += payload.get('dz', 0) * 0.01
            cam.distance = max(0.5, min(20.0, cam.distance))
        elif cam_type == 'set_distance':
            cam.distance = payload.get('distance', 3.0)
        elif cam_type == 'pan':
            right = np.array([np.sin(a), -np.cos(a), 0])
            up = np.array([-np.cos(a) * np.sin(e), -np.sin(a) * np.sin(e), np.cos(e)])
            scale = cam.distance * 0.002
            cam.lookat -= (payload.get('dx', 0) * right - payload.get('dy', 0) * up) * scale

    elif msg_type == 'camera_follow':
        payload = data.get('data', {})
        camera_follow = payload.get('follow', True)

    elif msg_type == 'toggle_target_visibility':
        payload = data.get('data', {})
        t_target_visible = payload.get('visible', True)

    elif msg_type == 'teleop_action':
        payload = data.get('data', {})
        # Accept both old format (dx/dy/dz) and new format (vx/vy/vz)
        vx = float(payload.get('vx', payload.get('dx', 0.0)))
        vy = float(payload.get('vy', payload.get('dy', 0.0)))
        vz = float(payload.get('vz', payload.get('dz', 0.0)))
        timestamp = time.time()

        if current_robot in ("ur5", "ur5_t_push") and _is_jogging_paused():
            return

        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                # Check if all velocities are zero - if so, stop any active jogging
                if abs(vx) < 0.001 and abs(vy) < 0.001 and abs(vz) < 0.001:
                    stop_jog = getattr(env, "stop_jog", None)
                    if callable(stop_jog):
                        stop_jog()
                else:
                    # Apply teleop movement
                    target = env.get_target()
                    env.set_target(
                        target[0] + vx,
                        target[1] + vy,
                        target[2] + vz,
                        update_joint_targets=True
                    )
                updated_target = env.get_target()
            else:
                updated_target = None

        with teleop_lock:
            # Update UR5 Cartesian velocities, preserve other fields
            last_teleop_action["vx"] = vx
            last_teleop_action["vy"] = vy
            last_teleop_action["vz"] = vz
            last_teleop_action["vyaw"] = 0.0
            # Clear jogging velocities when using teleop (jogging and teleop are mutually exclusive)
            for j in range(1, 7):
                last_teleop_action[f"j{j}"] = 0.0
            last_teleop_action["vrx"] = 0.0
            last_teleop_action["vry"] = 0.0
            last_teleop_action["vrz"] = 0.0

        broadcast_to_external_clients(
            "action_update",
            {
                "robot": current_robot,
                "scene": getattr(env, "scene_name", None) if env is not None else None,
                "delta": {"vx": vx, "vy": vy, "vz": vz},
                "ee_target": updated_target.tolist() if updated_target is not None else None,
                "timestamp": timestamp,
            }
        )

    elif msg_type == 'episode_control':
        payload = data.get('data', {})
        action = payload.get('action')
        _signal_episode_control(action)

    # UR5-specific messages
    elif msg_type == 'arm_target':
        payload = data.get('data', {})
        x = payload.get('x', 0.4)
        y = payload.get('y', 0.0)
        z = payload.get('z', 0.6)
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                env.set_target(x, y, z)

    elif msg_type == 'gripper':
        payload = data.get('data', {})
        action = payload.get('action', 'open')
        # Robotiq 2F-85: 0 = open, 255 = closed (official MuJoCo Menagerie)
        if action == 'open':
            value = 0  # 0 = fully open
        elif action == 'close':
            value = 255  # 255 = fully closed
        else:
            value = payload.get('value', 128)
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                env.set_gripper(value)
        # Update teleop_action with gripper value
        with teleop_lock:
            last_teleop_action["gripper"] = float(value)

    elif msg_type == 'control_mode':
        payload = data.get('data', {})
        mode = payload.get('mode', 'ik')
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                env.set_control_mode(mode)

    elif msg_type == 'joint_positions':
        payload = data.get('data', {})
        positions = payload.get('positions', [])
        if len(positions) == 6:
            with mujoco_lock:
                if env is not None and current_robot in ("ur5", "ur5_t_push"):
                    env.set_joint_positions(positions)

    elif msg_type == 'arm_orientation':
        payload = data.get('data', {})
        roll = payload.get('roll', 0.0)
        pitch = payload.get('pitch', np.pi/2)
        yaw = payload.get('yaw', 0.0)
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                env.set_target_orientation(roll, pitch, yaw)

    elif msg_type == 'use_orientation':
        payload = data.get('data', {})
        use = payload.get('enabled', True)
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push"):
                env.set_use_orientation(use)

    elif msg_type == 'perception_update':
        # Handle perception updates from training system
        payload = data.get('data', {})
        poses = payload.get('poses', [])
        timestamp = payload.get('timestamp', time.time())
        with perception_lock:
            perception_state["poses"] = poses
            perception_state["last_update"] = timestamp
        # print(f"[Perception] Received {len(poses)} pose(s) at {timestamp:.2f}")

    elif msg_type == 'set_nova_mode':
        payload = data.get('data', {})
        enabled = bool(payload.get('enabled', False))
        with mujoco_lock:
            if env is not None and current_robot in ("ur5", "ur5_t_push") and hasattr(env, 'set_nova_enabled'):
                env.set_nova_enabled(enabled)

    # Nova jogging commands
    elif msg_type == 'homing':
        with mujoco_lock:
            if env is not None:
                home_pose = getattr(env, "home_pose", None)
                set_joints = getattr(env, "set_joint_positions", None)
                if home_pose is not None and callable(set_joints):
                    if hasattr(home_pose, "tolist"):
                        positions = home_pose.tolist()
                    else:
                        positions = list(home_pose)
                    set_joints(positions)


def _serialize_value(value):
    """Recursively convert numpy data to JSON-friendly types."""
    if isinstance(value, np.ndarray):
        return _serialize_value(value.tolist())
    if isinstance(value, dict):
        return {str(k): _serialize_value(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_serialize_value(v) for v in value]
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        return float(value)
    return value


def _serialize_space(space):
    if hasattr(space, "low") and hasattr(space, "high"):
        import math
        # Convert numpy arrays to lists, replacing inf with null for valid JSON
        low_list = space.low.tolist()
        high_list = space.high.tolist()

        # Replace infinity values with null (JSON doesn't support Infinity)
        low_json = [None if math.isinf(v) else v for v in low_list]
        high_json = [None if math.isinf(v) else v for v in high_list]

        return {
            "type": "box",
            "shape": list(space.shape),
            "low": low_json,
            "high": high_json,
            "dtype": str(space.dtype),
        }
    return {"type": "unknown"}


def _resolve_robot_scene(robot, scene):
    if robot == "ur5_t_push" and scene is None:
        return "ur5", "scene_t_push"
    if robot == "ur5" and scene is None:
        return "ur5", "scene"
    return robot, scene


def _create_env(robot, scene):
    if robot == "g1":
        return G1Env(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT)
    if robot == "spot":
        spot_dir = os.path.join(_nova_sim_dir, 'robots', 'spot')
        sys.path.insert(0, spot_dir)
        from spot_env import SpotEnv
        sys.path.pop(0)
        return SpotEnv(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT, controller_type='pympc')
    if robot == "ur5":
        ur5_dir = os.path.join(_nova_sim_dir, 'robots', 'ur5')
        sys.path.insert(0, ur5_dir)
        from ur5_env import UR5Env
        sys.path.pop(0)
        # scene is already resolved by _resolve_robot_scene (e.g., "scene_t_push")
        if scene:
            return UR5Env(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT, scene_name=scene)
        return UR5Env(render_mode="rgb_array", width=RENDER_WIDTH, height=RENDER_HEIGHT)
    raise ValueError(f"Unsupported robot for gym: {robot}")


@sock.route(f'{API_PREFIX}/ws')
def websocket_handler(ws):
    """Handle WebSocket connections."""
    print('WebSocket client connected', flush=True)

    # Register client (track all clients, even before they identify)
    with ws_clients_lock:
        ws_clients.add(ws)

    # Register in metadata as unidentified
    _register_external_client(ws)
    broadcast_connected_clients_status()

    try:
        while True:
            message = ws.receive()
            if message is None:
                break
            try:
                data = json.loads(message)
                handle_ws_message(ws, data)
            except json.JSONDecodeError:
                print(f"Invalid JSON received: {message}")
            except Exception as e:
                print(f"Error handling message: {e}")
    except:
        pass
    finally:
        # Unregister client
        with ws_clients_lock:
            ws_clients.discard(ws)
        with external_ws_clients_lock:
            was_external = ws in external_ws_clients
            external_ws_clients.discard(ws)
        # Always unregister from metadata (all clients are now tracked, not just external ones)
        _unregister_external_client(ws)
        broadcast_state()
        broadcast_connected_clients_status()
        print('WebSocket client disconnected')


# Serve UI at /nova-sim (no redirect)
@app.route('/nova-sim')
@app.route('/nova-sim/')
def nova_sim_ui():
    return index()


# Redirect root to /nova-sim
@app.route('/')
def root_redirect():
    return render_template_string("""
    <!DOCTYPE html>
    <html>
    <head>
        <meta http-equiv="refresh" content="0; url='/nova-sim'" />
    </head>
    <body>
        <p>Redirecting to <a href="/nova-sim">/nova-sim</a>...</p>
    </body>
    </html>
    """)


@app.route(API_PREFIX)
@app.route(f'{API_PREFIX}/')
def index():
    # Load HTML from frontend/index.html
    html_path = Path(_nova_sim_dir) / 'frontend' / 'index.html'
    try:
        with open(html_path, 'r', encoding='utf-8') as f:
            html_template = f.read()
    except FileNotFoundError:
        return "Error: frontend/index.html not found", 500

    # Replace template variables (Python string concatenation syntax from old inline HTML)
    html_content = html_template.replace('""" + API_PREFIX + """', API_PREFIX)

    response = make_response(html_content)
    # Prevent browser caching of the HTML/JavaScript
    response.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = '0'
    return response


@app.route(f'{API_PREFIX}/video_feed')
def video_feed():
    return Response(generate_frames(),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


@app.route(f'{API_PREFIX}/cameras', methods=['POST'])
def add_camera():
    """Add an overlay camera for the current robot/scene."""
    payload = request.get_json(silent=True) or {}
    replace_existing = bool(payload.get("replace", False))

    try:
        config = _build_dynamic_camera_config(payload)
    except ValueError as exc:
        return jsonify({"error": str(exc)}), 400

    with mujoco_lock:
        if env is None:
            return jsonify({"error": "no environment initialized"}), 503

        scene_name = getattr(env, "scene_name", None) or current_scene
        robot_type = current_robot

        # Validate site references if provided
        follow_site = config.get("follow_site")
        look_target_site = config.get("look_target_site")
        if follow_site:
            site_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_SITE, follow_site)
            if site_id < 0:
                return jsonify({"error": f"follow_site '{follow_site}' not found"}), 400
        if look_target_site:
            target_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_SITE, look_target_site)
            if target_id < 0:
                return jsonify({"error": f"look_target_site '{look_target_site}' not found"}), 400

        temp_cam = _make_mjv_camera(config)
        feed_intrinsics = _camera_intrinsics(
            OVERLAY_RENDER_WIDTH,
            OVERLAY_RENDER_HEIGHT,
            _camera_fovy_degrees(temp_cam or cam, env),
        )
        camera_payload = {
            "name": config["name"],
            "label": config.get("label", config["name"]),
            "url": f"{API_PREFIX}/camera/{config['name']}/video_feed",
            "pose": _camera_pose(temp_cam),
            "intrinsics": feed_intrinsics,
        }

    with overlay_camera_lock:
        exists = config["name"] in overlay_camera_states
    if exists and not replace_existing:
        return jsonify({"error": "camera already exists", "name": config["name"]}), 409

    _store_dynamic_camera_config(robot_type, scene_name, config)

    state_payload = _build_overlay_camera_state(config, create_renderer=False)

    # Store reference to old renderer to close after state update (prevents race condition)
    old_renderer = None
    with overlay_camera_lock:
        if config["name"] in overlay_camera_states:
            old_renderer = overlay_camera_states[config["name"]].get("renderer")
        # Update state first before closing old renderer to prevent use-after-free
        overlay_camera_states[config["name"]] = state_payload
    with overlay_frame_lock:
        overlay_frames[config["name"]] = None

    # Close old renderer after state is updated (prevents simulation loop from accessing freed renderer)
    if old_renderer:
        try:
            old_renderer.close()
        except Exception:
            pass

    broadcast_camera_event(
        "added",
        camera_payload,
        {"robot": robot_type, "scene": scene_name},
    )
    return jsonify({
        "status": "ok",
        "camera": camera_payload,
        "scope": {"robot": robot_type, "scene": scene_name},
    })


@app.route(f'{API_PREFIX}/camera/<name>/video_feed')
def camera_feed(name):
    if name == "main":
        return Response(generate_frames(),
                        mimetype='multipart/x-mixed-replace; boundary=frame')

    with overlay_frame_lock:
        if name not in overlay_frames:
            return jsonify({"error": "Camera feed not available"}), 404

    return Response(generate_overlay_frames(name),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


def capture_depth_snapshot(camera_name: str) -> Optional[bytes]:
    """Capture a depth image from the specified camera and return as PNG bytes.

    Args:
        camera_name: Name of the camera ("main" or overlay camera name)

    Returns:
        PNG-encoded depth image as bytes, or None if capture fails
    """
    global env, cam, overlay_camera_states

    if env is None:
        return None

    # Determine which camera to use
    camera_obj = None
    if camera_name == "main":
        camera_obj = cam
    else:
        with overlay_camera_lock:
            if camera_name in overlay_camera_states:
                camera_obj = overlay_camera_states[camera_name].get("camera")

    if camera_obj is None:
        return None

    # Create a temporary renderer for depth capture
    with mujoco_lock:
        try:
            if camera_name == "main":
                render_height = env.height
                render_width = env.width
            else:
                render_height = OVERLAY_RENDER_HEIGHT
                render_width = OVERLAY_RENDER_WIDTH

            depth_renderer = mujoco.Renderer(
                env.model,
                height=render_height,
                width=render_width
            )

            # Enable depth rendering
            depth_renderer.enable_depth_rendering()

            # Update scene and render depth
            depth_renderer.update_scene(env.data, camera=camera_obj)
            depth_array = depth_renderer.render()  # Returns (height, width) float32

            # Close the temporary renderer
            depth_renderer.close()

            # Convert depth to 16-bit PNG in millimeters (metric depth).
            # MuJoCo renderer returns depth in meters; preserve metric scale.
            max_depth_m = 10.0
            try:
                vis_map = getattr(env.model, "vis", None)
                vis_map = getattr(vis_map, "map", None) if vis_map is not None else None
                zfar = getattr(vis_map, "zfar", None) if vis_map is not None else None
                if zfar is not None:
                    max_depth_m = float(zfar)
            except Exception:
                pass
            # uint16 in mm can represent up to 65.535 m
            max_depth_m = min(max_depth_m, 65.535)
            depth_m = np.nan_to_num(depth_array, nan=0.0, posinf=max_depth_m, neginf=0.0)
            depth_m = np.clip(depth_m, 0.0, max_depth_m)
            depth_16bit = (depth_m * 1000.0).round().astype(np.uint16)

            # Encode as PNG
            ret, buffer = cv2.imencode('.png', depth_16bit)
            if not ret:
                return None

            return buffer.tobytes()

        except Exception as e:
            print(f"Error capturing depth snapshot: {e}")
            traceback.print_exc()
            return None


@app.route(f'{API_PREFIX}/camera/<name>/depth_snapshot')
def camera_depth_snapshot(name):
    """Endpoint to capture a single depth image as PNG."""
    depth_png = capture_depth_snapshot(name)

    if depth_png is None:
        return jsonify({"error": "Failed to capture depth image"}), 500

    response = make_response(depth_png)
    response.headers['Content-Type'] = 'image/png'
    response.headers['Content-Disposition'] = f'inline; filename="depth_{name}.png"'
    return response


@app.route(f'{API_PREFIX}/metadata')
def metadata():
    robots_meta = {}
    for robot, scenes in ROBOT_SCENES.items():
        robots_meta[robot] = {
            "label": ROBOT_LABELS.get(robot, robot),
            "scenes": scenes,
        }
    current_selection = {
        "robot": current_robot,
        "scene": current_scene,
    }
    with mujoco_lock:
        if env is not None and hasattr(env, "home_pose"):
            try:
                home_pose = env.home_pose
                if home_pose is not None:
                    pose_list = (
                        home_pose.tolist()
                        if hasattr(home_pose, "tolist")
                        else list(home_pose)
                    )
                    current_selection["home_pose"] = [float(val) for val in pose_list]
            except Exception:
                pass
    response = {
        "robots": robots_meta,
        "actions": AVAILABLE_ACTIONS,
        "nova_api": {
            "preconfigured": _nova_credentials_present(),
        },
    }
    return jsonify(response)


@app.route(f'{API_PREFIX}/nova_config', methods=['POST'])
def update_nova_config():
    payload = request.get_json(silent=True) or {}
    field_map = {
        "instance_url": "NOVA_INSTANCE_URL",
        "access_token": "NOVA_ACCESS_TOKEN",
        "cell_id": "NOVA_CELL_ID",
        "controller_id": "NOVA_CONTROLLER_ID",
        "motion_group_id": "NOVA_MOTION_GROUP_ID",
        "motion_group_model": "NOVA_MOTION_GROUP_MODEL",
        "tcp_name": "NOVA_TCP_NAME",
        "response_rate_ms": "NOVA_RESPONSE_RATE_MS",
    }
    updated = {}
    for key, env_key in field_map.items():
        if key not in payload:
            continue
        value = payload.get(key)
        if value is None or (isinstance(value, str) and value.strip() == ""):
            os.environ.pop(env_key, None)
            updated[env_key] = None
        else:
            os.environ[env_key] = str(value).strip()
            updated[env_key] = os.environ[env_key]

    if not _nova_credentials_present():
        return jsonify({
            "ok": False,
            "error": "Missing required Nova credentials",
            "updated": updated,
            "preconfigured": _nova_credentials_present(),
        }), 400

    use_state = _parse_bool(os.environ.get("NOVA_UR5_USE_STATE_STREAM") or os.environ.get("NOVA_USE_STATE_STREAM") or "true")
    use_ik = _parse_bool(os.environ.get("NOVA_UR5_USE_IK") or os.environ.get("NOVA_USE_IK") or "false")
    use_jogging = _parse_bool(os.environ.get("NOVA_UR5_USE_JOGGING") or os.environ.get("NOVA_USE_JOGGING") or "true")

    nova_api_config = {
        "instance_url": os.getenv("NOVA_INSTANCE_URL") or os.getenv("NOVA_API"),
        "access_token": os.getenv("NOVA_ACCESS_TOKEN"),
        "cell_id": os.getenv("NOVA_CELL_ID", "cell"),
        "controller_id": os.getenv("NOVA_CONTROLLER_ID"),
        "motion_group_id": os.getenv("NOVA_MOTION_GROUP_ID"),
        "motion_group_model": os.getenv("NOVA_MOTION_GROUP_MODEL"),
        "tcp_name": os.getenv("NOVA_TCP_NAME"),
        "response_rate_ms": int(os.getenv("NOVA_RESPONSE_RATE_MS", "200")),
        "use_state_stream": use_state,
        "use_ik": use_ik,
        "use_jogging": use_jogging,
    }

    reconfigured = False
    with mujoco_lock:
        if env is not None and current_robot in ("ur5", "ur5_t_push") and hasattr(env, "set_nova_enabled"):
            try:
                if hasattr(env, "_nova_api_config"):
                    env._nova_api_config = dict(nova_api_config)
                reconfigured = bool(env.set_nova_enabled(True))
            except Exception:
                traceback.print_exc()
                return jsonify({
                    "ok": False,
                    "error": "Failed to reconfigure Nova integration",
                    "updated": updated,
                    "preconfigured": _nova_credentials_present(),
                }), 500

    return jsonify({
        "ok": True,
        "updated": updated,
        "preconfigured": _nova_credentials_present(),
        "reconfigured": reconfigured,
    })


def _env_request(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception:
            traceback.print_exc()
            raise
    return wrapper


@app.route(f'{API_PREFIX}/env')
@_env_request
def get_env_info():
    """Get static information about the current environment."""
    with mujoco_lock:
        if env is None:
            return jsonify({"error": "no environment initialized"}), 503

        # Get action and observation spaces
        action_space = _serialize_space(env.action_space)
        observation_space = _serialize_space(env.observation_space)

        # Get static environment properties
        has_gripper = getattr(env, 'has_gripper', False)
        scene_name = getattr(env, 'scene_name', None)
        control_mode = getattr(env, 'get_control_mode', lambda: None)()

        # Build camera feed URLs with pose + intrinsics
        camera_feeds = []
        main_fovy = _camera_fovy_degrees(cam, env)
        main_intrinsics = _camera_intrinsics(env.width, env.height, main_fovy)
        camera_feeds.append({
            "name": "main",
            "label": "Main Camera",
            "url": f"{API_PREFIX}/video_feed",
            "pose": _camera_pose(cam),
            "intrinsics": main_intrinsics,
        })

        # Check for overlay cameras based on scene
        configs = (scene_name and OVERLAY_CAMERA_PRESETS.get(scene_name)) or OVERLAY_CAMERA_PRESETS.get(current_robot) or []
        dynamic_configs = _get_dynamic_camera_configs(current_robot, scene_name)
        if dynamic_configs:
            merged_configs: dict[str, dict[str, Any]] = {}
            for feed in configs:
                merged_configs[feed.get("name", "aux")] = feed
            for feed in dynamic_configs:
                merged_configs[feed.get("name", "aux")] = feed
            configs = list(merged_configs.values())
        for feed in configs:
            feed_name = feed.get('name', 'aux')
            with overlay_camera_lock:
                cam_state = overlay_camera_states.get(feed_name, {}).copy()
            cam_obj = cam_state.get("camera")
            feed_intrinsics = _camera_intrinsics(
                OVERLAY_RENDER_WIDTH,
                OVERLAY_RENDER_HEIGHT,
                _camera_fovy_degrees(cam_obj or cam, env),
            )
            feed_payload = {
                "name": feed_name,
                "label": feed.get('label', feed_name),
                "url": f"{API_PREFIX}/camera/{feed_name}/video_feed"
            }
            if cam_obj is not None:
                feed_payload["pose"] = _camera_pose(cam_obj)
            feed_payload["intrinsics"] = feed_intrinsics
            camera_feeds.append({
                **feed_payload
            })

        # Get home pose for current robot
        home_pose_list = None
        if hasattr(env, "home_pose"):
            try:
                home_pose = env.home_pose
                if home_pose is not None:
                    pose_list = (
                        home_pose.tolist()
                        if hasattr(home_pose, "tolist")
                        else list(home_pose)
                    )
                    home_pose_list = [float(val) for val in pose_list]
            except Exception:
                pass

        response = {
            "robot": current_robot,
            "scene": scene_name or current_scene,
            "has_gripper": has_gripper,
            "control_mode": control_mode,
            "action_space": action_space,
            "observation_space": observation_space,
            "camera_feeds": camera_feeds
        }

        # Add home_pose if available
        if home_pose_list is not None:
            response["home_pose"] = home_pose_list

        return jsonify(response)


@app.route(f'{API_PREFIX}/homing', methods=['GET'])
@_env_request
def home_blocking():
    timeout_s = float(request.args.get("timeout_s", 30.0))
    tolerance = float(request.args.get("tolerance", 0.01))
    poll_interval = float(request.args.get("poll_interval_s", 0.1))

    if timeout_s <= 0:
        return jsonify({"error": "timeout_s must be > 0"}), 400
    if tolerance <= 0:
        return jsonify({"error": "tolerance must be > 0"}), 400
    if poll_interval <= 0:
        return jsonify({"error": "poll_interval_s must be > 0"}), 400

    with mujoco_lock:
        if env is None:
            return jsonify({"error": "no environment initialized"}), 503

    if current_robot not in ("ur5", "ur5_t_push"):
        with mujoco_lock:
            home_pose = getattr(env, "home_pose", None)
            set_joints = getattr(env, "set_joint_positions", None)
            if home_pose is None or not callable(set_joints):
                return jsonify({"error": "homing not supported for current robot"}), 400
            positions = home_pose.tolist() if hasattr(home_pose, "tolist") else list(home_pose)
            set_joints(positions)
        return jsonify({"status": "completed", "robot": current_robot, "homed": True})

    with homing_lock:
        homing_active = homing_state.get("active")
    if homing_active:
        active_start_time = time.time()
        while True:
            with homing_lock:
                if not homing_state.get("active"):
                    break
            if time.time() - active_start_time >= timeout_s:
                return jsonify({"status": "timeout", "robot": current_robot, "homed": False}), 200
            time.sleep(min(poll_interval, 0.5))

    with mujoco_lock, homing_lock:
        home_pose = getattr(env, "home_pose", None)
        get_joints = getattr(env, "get_joint_positions", None)
        start_multi_joint_jog = getattr(env, "start_multi_joint_jog", None)
        stop_jog = getattr(env, "stop_jog", None)
        if home_pose is None or not callable(get_joints) or not callable(start_multi_joint_jog) or not callable(stop_jog):
            return jsonify({"error": "Nova homing not supported for current robot"}), 400
        homing_state["active"] = True

    previous_pause = _set_jogging_paused(True)
    start_time = time.time()
    status = "running"
    result_snapshot = {}

    try:
        with mujoco_lock:
            stop_jog = getattr(env, "stop_jog", None)
            if callable(stop_jog):
                stop_jog()

        while True:
            with mujoco_lock, homing_lock:
                result_snapshot = _step_multi_axis_homing_locked(tolerance=tolerance)

            if result_snapshot.get("status") == "done":
                status = "completed"
                break

            if time.time() - start_time >= timeout_s:
                status = "timeout"
                break

            time.sleep(poll_interval)
    finally:
        with mujoco_lock, homing_lock:
            stop_jog = getattr(env, "stop_jog", None)
            if callable(stop_jog) and homing_state.get("is_jogging", False):
                stop_jog()
            homing_state["active"] = False
            homing_state["is_jogging"] = False
            homing_state.pop("previous_errors", None)
            homing_state.pop("current_velocities", None)
        _set_jogging_paused(previous_pause)

    with mujoco_lock:
        home_pose = getattr(env, "home_pose", None)
        get_joints = getattr(env, "get_joint_positions", None)
        final_positions = get_joints() if callable(get_joints) else []

    home_pose_list = (
        home_pose.tolist()
        if home_pose is not None and hasattr(home_pose, "tolist")
        else (list(home_pose) if home_pose is not None else None)
    )

    final_errors = None
    max_abs_error = None
    if home_pose_list is not None and final_positions is not None and len(final_positions) > 0:
        final_errors = [float(t - c) for t, c in zip(home_pose_list, final_positions)]
        max_abs_error = max(abs(e) for e in final_errors)

    final_positions_list = [float(v) for v in final_positions] if final_positions is not None and len(final_positions) > 0 else None
    response = {
        "status": status,
        "robot": current_robot,
        "homed": status == "completed",
        "duration_s": time.time() - start_time,
        "home_pose": home_pose_list,
        "final_positions": final_positions_list,
        "final_errors": final_errors,
        "max_abs_error": max_abs_error,
        "jogging_paused": _is_jogging_paused(),
    }
    return jsonify(response)




if __name__ == '__main__':
    port = int(os.environ.get("PORT", "3004"))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)