#!/usr/bin/env bash
# scripts/sim/run_sim.sh
#
# Worker script invoked by `agenticros up sim-amr` (and later `sim-arm`). Sources
# ROS 2 + the AgenticROS overlay, runs `ros2 launch agenticros_sim ...`, writes
# its PID into /tmp/agenticros-sim.pid so `agenticros down` can stop it, and
# tees output to /tmp/agenticros-sim.log so `agenticros logs sim` works.
#
# Usage:
#   run_sim.sh --robot amr [--namespace sim_robot] [--rviz] [--no-gui] [--nav2] [--real-camera]
#   run_sim.sh --robot arm [--namespace sim_robot] [--rviz] [--no-gui] [--moveit]
#
# Flags:
#   --robot amr|arm        Which sim launch to start (default: amr).
#   --namespace <ns>       Robot namespace; exported as AGENTICROS_ROBOT_NAMESPACE.
#   --rviz                 Bring up RViz alongside Gazebo.
#   --no-gui               Run gz-sim headless (CI / docker).
#   --nav2                 AMR only: also launch Nav2 (map + AMCL + navigation).
#   --real-camera          AMR only: live RealSense eyes on the sim AMR body (implies overlay).
#   --moveit               Arm only: also launch MoveIt2 move_group + trajectory bridge.
#   --ros-distro <distro>  Override ROS 2 distro (auto-detect by default).
#   --colcon-ws <path>     Override the colcon workspace (default: <repo>/ros2_ws).
#   --help                 Print this help and exit.
#
# Exit codes:
#   0   sim exited cleanly
#   1   missing dependency (gz, ros2, ros_gz_*, our launch file)
#   2   bad CLI usage
#   3   sim crashed at runtime — check /tmp/agenticros-sim.log

# Strict mode, but `set -u` clashes with ROS setup scripts that reference
# AMENT_TRACE_SETUP_FILES etc. We toggle nounset off around each `source`.
set -eo pipefail
shopt -s expand_aliases || true

# ---------- args ----------
ROBOT="amr"
NAMESPACE=""
USE_RVIZ="false"
GUI="true"
USE_NAV2="false"
USE_MOVEIT="false"
REAL_CAMERA="false"
ROS_DISTRO_OVERRIDE=""
COLCON_WS=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --robot)        ROBOT="$2"; shift 2 ;;
    --namespace)    NAMESPACE="$2"; shift 2 ;;
    --rviz)         USE_RVIZ="true"; shift ;;
    --no-gui)       GUI="false"; shift ;;
    --nav2)         USE_NAV2="true"; shift ;;
    --real-camera)  REAL_CAMERA="true"; shift ;;
    --moveit)       USE_MOVEIT="true"; shift ;;
    --ros-distro)   ROS_DISTRO_OVERRIDE="$2"; shift 2 ;;
    --colcon-ws)    COLCON_WS="$2"; shift 2 ;;
    -h|--help)
      sed -n '2,30p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
      exit 0 ;;
    *) echo "Unknown arg: $1" >&2; exit 2 ;;
  esac
done

if [[ "$ROBOT" != "amr" && "$ROBOT" != "arm" ]]; then
  echo "--robot must be 'amr' or 'arm' (got '$ROBOT')" >&2
  exit 2
fi

if [[ "$USE_NAV2" == "true" && "$ROBOT" != "amr" ]]; then
  echo "--nav2 is only supported with --robot amr" >&2
  exit 2
fi

if [[ "$REAL_CAMERA" == "true" && "$ROBOT" != "amr" ]]; then
  echo "--real-camera is only supported with --robot amr" >&2
  exit 2
fi

if [[ "$USE_MOVEIT" == "true" && "$ROBOT" != "arm" ]]; then
  echo "--moveit is only supported with --robot arm" >&2
  exit 2
fi

# ---------- paths ----------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
COLCON_WS="${COLCON_WS:-$REPO_ROOT/ros2_ws}"

LOG_FILE="/tmp/agenticros-sim.log"
PID_FILE="/tmp/agenticros-sim.pid"

log() { printf "\033[36m[run_sim]\033[0m %s\n" "$*"; }
err() { printf "\033[31m[run_sim]\033[0m %s\n" "$*" >&2; }

# ---------- ROS 2 distro detection ----------
if [[ -n "$ROS_DISTRO_OVERRIDE" ]]; then
  ROS_DISTRO="$ROS_DISTRO_OVERRIDE"
elif [[ -n "${ROS_DISTRO:-}" ]]; then
  : # already sourced
else
  for d in humble jazzy iron rolling; do
    if [[ -f "/opt/ros/$d/setup.bash" ]]; then ROS_DISTRO="$d"; break; fi
  done
fi
if [[ -z "${ROS_DISTRO:-}" ]] || [[ ! -f "/opt/ros/$ROS_DISTRO/setup.bash" ]]; then
  err "ROS 2 not found under /opt/ros/. Install Humble or Jazzy first."
  exit 1
fi

log "ROS_DISTRO=$ROS_DISTRO"
log "COLCON_WS=$COLCON_WS"
log "robot=$ROBOT  namespace=${NAMESPACE:-<none>}  rviz=$USE_RVIZ  gui=$GUI  nav2=$USE_NAV2  moveit=$USE_MOVEIT  real_camera=$REAL_CAMERA"

# shellcheck disable=SC1090
source "/opt/ros/$ROS_DISTRO/setup.bash"

# ---------- Build agenticros_sim if not yet built ----------
NEED_BUILD=0
if [[ ! -f "$COLCON_WS/install/agenticros_sim/share/agenticros_sim/launch/sim_amr.launch.py" ]]; then
  NEED_BUILD=1
fi
if [[ "$USE_MOVEIT" == "true" ]]; then
  if [[ ! -f "$COLCON_WS/install/agenticros_sim/share/agenticros_sim/launch/sim_arm_moveit.launch.py" ]] ||
     [[ ! -f "$COLCON_WS/install/agenticros_arm_moveit_config/share/agenticros_arm_moveit_config/launch/move_group.launch.py" ]]; then
    NEED_BUILD=1
  fi
fi
if [[ "$NEED_BUILD" -eq 1 ]]; then
  log "sim packages not built in $COLCON_WS — building now..."
  BUILD_PKGS=(agenticros_sim agenticros_msgs)
  if [[ "$USE_MOVEIT" == "true" ]]; then
    BUILD_PKGS+=(agenticros_arm_moveit_config)
  fi
  (cd "$COLCON_WS" && colcon build --symlink-install --packages-select "${BUILD_PKGS[@]}")
fi

# shellcheck disable=SC1090
source "$COLCON_WS/install/setup.bash"

# ---------- Dependency checks ----------
for bin in gz ros2 rviz2; do
  if ! command -v "$bin" >/dev/null; then
    if [[ "$bin" == "rviz2" ]] && [[ "$USE_RVIZ" != "true" ]]; then
      continue
    fi
    err "Required binary '$bin' not found on PATH."
    if [[ "$bin" == "gz" ]]; then
      err "Gazebo Sim is required for sim-amr. On ROS 2 ${ROS_DISTRO}:"
      err "  sudo apt install ros-${ROS_DISTRO}-ros-gz"
      err "Then re-run: ./agenticros up sim-amr --real-camera"
    fi
    exit 1
  fi
done

# ---------- Wire up environment ----------
if [[ -n "$NAMESPACE" ]]; then
  export AGENTICROS_ROBOT_NAMESPACE="$NAMESPACE"
fi

# Jetson rendering fix. On Tegra boards Mesa is picked first and tries to load
# nvidia-drm_dri.so which doesn't exist, so the gz GUI viewport comes up solid
# white. We:
#   1. Point libglvnd at the NVIDIA EGL vendor (works for some Jetson L4T images)
#   2. As a fallback, allow AGENTICROS_GZ_SOFTWARE_RENDER=1 to force llvmpipe -
#      slow (~5 fps) but actually renders the world so the demo is viewable.
# Honor AGENTICROS_GZ_NO_TWEAKS=1 to skip both (useful on x86/laptops).
if [[ "$GUI" == "true" ]] && [[ -z "${AGENTICROS_GZ_NO_TWEAKS:-}" ]]; then
  if [[ -f /usr/lib/aarch64-linux-gnu/tegra-egl/libEGL_nvidia.so.0 ]]; then
    log "Jetson detected: forcing NVIDIA EGL/GL vendor"
    export __GLX_VENDOR_LIBRARY_NAME="${__GLX_VENDOR_LIBRARY_NAME:-nvidia}"
    export __EGL_VENDOR_LIBRARY_FILENAMES="${__EGL_VENDOR_LIBRARY_FILENAMES:-/usr/share/glvnd/egl_vendor.d/10_nvidia.json}"
    export LD_LIBRARY_PATH="/usr/lib/aarch64-linux-gnu/tegra-egl:/usr/lib/aarch64-linux-gnu/tegra:${LD_LIBRARY_PATH:-}"
  fi
  if [[ -n "${AGENTICROS_GZ_SOFTWARE_RENDER:-}" ]]; then
    log "AGENTICROS_GZ_SOFTWARE_RENDER set - forcing Mesa llvmpipe software renderer"
    export LIBGL_ALWAYS_SOFTWARE=1
    export GALLIUM_DRIVER=llvmpipe
    export MESA_GL_VERSION_OVERRIDE=4.5
    export OGRE_RTT_MODE=Copy
  fi
fi

LAUNCH_ARGS=(
  "use_rviz:=$USE_RVIZ"
  "gui:=$GUI"
)

if [[ "$REAL_CAMERA" == "true" ]]; then
  LAUNCH_ARGS+=("real_camera:=true")
fi

# ---------- Pick launch file ----------
case "$ROBOT" in
  amr)
    if [[ "$USE_NAV2" == "true" ]]; then
      LAUNCH_FILE="sim_amr_nav2.launch.py"
    else
      LAUNCH_FILE="sim_amr.launch.py"
    fi
    ;;
  arm)
    if [[ "$USE_MOVEIT" == "true" ]]; then
      LAUNCH_FILE="sim_arm_moveit.launch.py"
    else
      LAUNCH_FILE="sim_arm.launch.py"
    fi
    ;;
esac

if [[ "$USE_NAV2" == "true" ]]; then
  if ! ros2 pkg prefix nav2_bringup >/dev/null 2>&1; then
    err "nav2_bringup not found. Install: sudo apt-get install -y ros-\$ROS_DISTRO-navigation2 ros-\$ROS_DISTRO-nav2-bringup"
    exit 1
  fi
fi

if [[ "$USE_MOVEIT" == "true" ]]; then
  if ! ros2 pkg prefix moveit_ros_move_group >/dev/null 2>&1; then
    err "moveit_ros_move_group not found. Install: sudo apt-get install -y ros-\$ROS_DISTRO-moveit"
    exit 1
  fi
  if ! ros2 pkg prefix agenticros_arm_moveit_config >/dev/null 2>&1; then
    err "agenticros_arm_moveit_config not found. Rebuild: colcon build --packages-select agenticros_sim agenticros_arm_moveit_config --symlink-install"
    exit 1
  fi
fi

if ! ros2 launch --help >/dev/null 2>&1; then
  err "ros2 launch not available — is ROS sourced properly?"
  exit 1
fi

# OpenClaw / MCP talk to this sim over rosbridge (ws://localhost:9090).
# local DDS via rclnodejs fails to rebuild on Node 22+ (Thor/Jetson), and
# without a working transport the chat agent falls back to bash + turtlesim.
start_sim_rosbridge() {
  local pidfile="/tmp/agenticros-rosbridge.pid"
  local logfile="/tmp/agenticros-rosbridge.log"
  if [[ -f "$pidfile" ]]; then
    local old
    old="$(tr -d '[:space:]' < "$pidfile" 2>/dev/null || true)"
    if [[ -n "$old" ]] && kill -0 "$old" 2>/dev/null; then
      log "rosbridge already running (pid $old) — OpenClaw uses ws://localhost:9090"
      return 0
    fi
  fi
  if ! ros2 pkg prefix rosbridge_server >/dev/null 2>&1; then
    err "rosbridge_server not found. OpenClaw cannot drive this sim without it."
    err "  sudo apt install -y ros-${ROS_DISTRO}-rosbridge-suite"
    return 0
  fi
  log "Starting rosbridge on ws://localhost:9090 (logs: $logfile)"
  nohup ros2 launch rosbridge_server rosbridge_websocket_launch.xml >>"$logfile" 2>&1 &
  echo $! > "$pidfile"
  disown || true
}
start_sim_rosbridge

log "Logging to $LOG_FILE"
log "ros2 launch agenticros_sim $LAUNCH_FILE ${LAUNCH_ARGS[*]}"
log "Press Ctrl+C to stop (or run \`agenticros down\` from another terminal)."

# Run in foreground so Ctrl+C in the parent shell still works, but also tee to
# the log so `agenticros logs sim -f` can follow concurrently. The PID we record
# is this script's own PID so `agenticros down` SIGTERMs the whole subtree.
echo "$$" > "$PID_FILE"

# Use `exec` so the ros2 process replaces our shell — that way signals from the
# CLI (or Ctrl+C) hit ros2 directly instead of bash.
exec 1> >(tee -a "$LOG_FILE") 2>&1
exec ros2 launch agenticros_sim "$LAUNCH_FILE" "${LAUNCH_ARGS[@]}"