#!/usr/bin/env bash
set -euo pipefail

# Unified launcher for TurtleBot3 office workflow:
# - Terminal A: simulation / bridge / (optional SLAM) / RViz
# - Terminal B: Nav2
# - Optional C: explore_lite
# - Optional websocket + gzweb-demo frontend
#
# Modes:
#   mapping       : online SLAM mapping
#   localization  : map-based localization/navigation
#
# Mapping mode interactive commands:
#   s : save map and switch to localization mode
#   q : stop all
#
# Examples:
#   ./run_all.sh mapping --explore --gzweb
#   ./run_all.sh localization --gzweb --map /abs/path/office_map.yaml

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODE="${1:-mapping}"
shift || true

ROBOT_MODEL="${ROBOT_MODEL:-waffle}"
USE_SIM_TIME="${USE_SIM_TIME:-true}"
GZ_PARTITION="${GZ_PARTITION:-tb3_office_gz}"
WORLD_PATH_DEFAULT="${ROOT}/turtlebot3_simulations/turtlebot3_gazebo/worlds/office_gz_dartsim.sdf"
MAP_YAML_DEFAULT="${ROOT}/turtlebot3_simulations/turtlebot3_gazebo/map/office_map.yaml"
MAP_BASENAME_DEFAULT="${ROOT}/turtlebot3_simulations/turtlebot3_gazebo/map/office_map"
SLEEP_NAV2="${SLEEP_NAV2:-18}"
EXPLORE_WAIT_TIMEOUT="${EXPLORE_WAIT_TIMEOUT:-420}"

WORLD_PATH="${WORLD_PATH:-$WORLD_PATH_DEFAULT}"
MAP_YAML="${NAV2_MAP_PATH:-$MAP_YAML_DEFAULT}"
MAP_BASENAME="${MAP_BASENAME:-$MAP_BASENAME_DEFAULT}"

WITH_EXPLORE=false
WITH_GZWEB=false
NO_RVIZ=false
# 与 --gzweb 同用：只起 gz websocket(9002) 等资源，不起 office 内 npm run dev；页面用镜像内 /opt/gzweb（foundation）
NO_GZWEB_VITE=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    --explore) WITH_EXPLORE=true; shift ;;
    --gzweb) WITH_GZWEB=true; shift ;;
    --no-gzweb-vite) NO_GZWEB_VITE=true; shift ;;
    --no-rviz) NO_RVIZ=true; shift ;;
    --map) MAP_YAML="$2"; shift 2 ;;
    --world) WORLD_PATH="$2"; shift 2 ;;
    --help|-h)
      cat <<'EOF'
Usage:
  ./run_all.sh mapping [--explore] [--gzweb] [--no-gzweb-vite] [--no-rviz] [--world PATH]
  ./run_all.sh localization [--gzweb] [--no-gzweb-vite] [--no-rviz] [--map PATH] [--world PATH]

  --no-gzweb-vite  需与 --gzweb 同时使用：不启动 workspace 内 gzweb-demo 的 Vite；浏览器用 /opt/gzweb（8000）。
EOF
      exit 0
      ;;
    *)
      echo "[run_all] unknown argument: $1" >&2
      exit 1
      ;;
  esac
done

if [[ "$MODE" != "mapping" && "$MODE" != "localization" ]]; then
  echo "[run_all] mode must be mapping or localization" >&2
  exit 1
fi

if [[ ! -f "$WORLD_PATH" ]]; then
  echo "[run_all] world not found: $WORLD_PATH" >&2
  exit 1
fi

if [[ "$MODE" == "localization" && ! -f "$MAP_YAML" ]]; then
  echo "[run_all] map yaml not found: $MAP_YAML" >&2
  exit 1
fi

set +u
source /opt/ros/humble/setup.bash
source "${ROOT}/install/setup.bash"
set -u
export TURTLEBOT3_MODEL="$ROBOT_MODEL"

LOG_DIR="${ROOT}/logs/run_all_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$LOG_DIR"
mkdir -p "$(dirname "$MAP_BASENAME")"

PIDS=()

start_bg() {
  local name="$1"; shift
  local logfile="${LOG_DIR}/${name}.log"
  echo "[run_all] starting ${name} ..."
  "$@" >"$logfile" 2>&1 &
  local pid=$!
  PIDS+=("$pid")
  echo "[run_all] ${name} pid=${pid}, log=${logfile}"
}

cleanup() {
  echo "[run_all] stopping all processes..."
  for pid in "${PIDS[@]:-}"; do
    kill "$pid" 2>/dev/null || true
  done
}
trap cleanup EXIT INT TERM

wait_for_topic() {
  local topic="$1"
  local timeout="${2:-60}"
  local i=0
  while (( i < timeout )); do
    if ros2 topic list 2>/dev/null | grep -Fxq "${topic}"; then
      return 0
    fi
    sleep 1
    ((i+=1))
  done
  return 1
}

wait_for_port() {
  local port="$1"
  local timeout="${2:-30}"
  local i=0
  while (( i < timeout )); do
    if ss -ltn 2>/dev/null | grep -q ":${port} "; then
      return 0
    fi
    sleep 1
    ((i+=1))
  done
  return 1
}

wait_for_explore_stop_log() {
  local logfile="$1"
  local timeout="${2:-900}"
  local i=0
  while (( i < timeout )); do
    if [[ -f "$logfile" ]] && grep -Eq "All frontiers traversed/tried out, stopping\\.|Exploration stopped\\." "$logfile"; then
      return 0
    fi
    sleep 1
    ((i+=1))
  done
  return 1
}

wait_for_explore_stop_log_with_timer() {
  local logfile="$1"
  local timeout="${2:-900}"
  local start_ts
  start_ts="$(date +%s)"
  local i=0
  while (( i < timeout )); do
    if [[ -f "$logfile" ]] && grep -Eq "All frontiers traversed/tried out, stopping\\.|Exploration stopped\\." "$logfile"; then
      return 0
    fi
    if (( i % 10 == 0 )); then
      local now_ts elapsed
      now_ts="$(date +%s)"
      elapsed=$(( now_ts - start_ts ))
      echo "[run_all] explore elapsed: ${elapsed}s"
    fi
    sleep 1
    ((i+=1))
  done
  return 1
}

kill_port_listeners() {
  local port="$1"
  local pids
  pids="$(ss -ltnp 2>/dev/null | grep ":${port} " | sed -E 's/.*pid=([0-9]+).*/\1/' | sort -u || true)"
  if [[ -n "${pids}" ]]; then
    echo "[run_all] port ${port} already in use, killing: ${pids}"
    kill ${pids} 2>/dev/null || true
    sleep 1
  fi
}

show_rviz_arg="show_rviz:=true"
if [[ "$NO_RVIZ" == true ]]; then
  show_rviz_arg="show_rviz:=false"
fi

if [[ "$MODE" == "mapping" ]]; then
  start_bg "A_sim" \
    ros2 launch turtlebot3_gazebo turtlebot3_office_gz.launch.py \
      robot_model:="$ROBOT_MODEL" \
      use_sim_time:="$USE_SIM_TIME" \
      start_slam:=true \
      "$show_rviz_arg" \
      gz_partition:="$GZ_PARTITION" \
      gz_world_name:=default \
      world:="$WORLD_PATH"
else
  start_bg "A_sim" \
    ros2 launch turtlebot3_gazebo turtlebot3_office_gz.launch.py \
      robot_model:="$ROBOT_MODEL" \
      use_sim_time:="$USE_SIM_TIME" \
      start_slam:=false \
      "$show_rviz_arg" \
      gz_partition:="$GZ_PARTITION" \
      gz_world_name:=default \
      world:="$WORLD_PATH"
fi

# --gzweb：gz launch(9002) + 可选 office 内 gzweb-demo 的 Vite（5173）。
# --gzweb --no-gzweb-vite：不起 office Vite；页面用 foundation 已在 /opt/gzweb:8000 起的实例（见 run_all_fix.sh）。
if [[ "$WITH_GZWEB" == true ]]; then
  export GZ_PARTITION
  OFFICE_MODEL_ROOT="${ROOT}/turtlebot3_simulations/turtlebot3_gazebo/models/turtlebot3_office"
  RESOURCE_PATHS=(
    "${HOME}/.gazebo/models"
    "${ROOT}/turtlebot3_simulations/turtlebot3_gazebo/models"
    "${OFFICE_MODEL_ROOT}"
  )
  for d in "${OFFICE_MODEL_ROOT}"/*/materials/scripts "${OFFICE_MODEL_ROOT}"/*/materials/textures; do
    [[ -d "${d}" ]] && RESOURCE_PATHS+=("${d}")
  done
  RESOURCE_PATH_JOINED="$(IFS=:; echo "${RESOURCE_PATHS[*]}")"
  export GZ_SIM_RESOURCE_PATH="${RESOURCE_PATH_JOINED}"
  # Keep compatibility with components resolving textures via legacy lookup vars.
  export GZ_FILE_PATH="${RESOURCE_PATH_JOINED}"
  export GAZEBO_RESOURCE_PATH="${RESOURCE_PATH_JOINED}"

  if [[ "$NO_GZWEB_VITE" != true ]]; then
    if [[ ! -d "${ROOT}/gzweb-demo/gzweb-demo/node_modules/gzweb" ]]; then
      echo "[run_all] gzweb demo dependencies missing, running npm install once..."
      bash -lc "cd '${ROOT}/gzweb-demo/gzweb-demo' && npm install"
    fi
  else
    echo "[run_all] --no-gzweb-vite: skipping office gzweb-demo npm install; use /opt/gzweb (port 8000) for UI."
  fi

  kill_port_listeners 9002
  kill_port_listeners 5173

  echo "[run_all] waiting for /clock before websocket..."
  if wait_for_topic "/clock" 25; then
    start_bg "D_websocket" gz launch "${ROOT}/websocket.sdf"
    if [[ "$NO_GZWEB_VITE" != true ]]; then
      start_bg "E_gzweb_frontend" bash -lc "cd '${ROOT}/gzweb-demo/gzweb-demo' && npm run dev:host"
      echo "[run_all] office gzweb-demo Vite: http://localhost:5173"
    fi
  else
    echo "[run_all] WARN: /clock not detected in time; starting websocket (and optional frontend) anyway."
    start_bg "D_websocket" gz launch "${ROOT}/websocket.sdf"
    if [[ "$NO_GZWEB_VITE" != true ]]; then
      start_bg "E_gzweb_frontend" bash -lc "cd '${ROOT}/gzweb-demo/gzweb-demo' && npm run dev:host"
    fi
  fi

  if wait_for_port 9002 12; then
    echo "[run_all] websocket server ready on ws://localhost:9002"
  else
    echo "[run_all] ERROR: websocket server failed on 9002; check ${LOG_DIR}/D_websocket.log"
  fi
fi

echo "[run_all] waiting ${SLEEP_NAV2}s before starting Nav2..."
sleep "$SLEEP_NAV2"

if [[ "$MODE" == "mapping" ]]; then
  start_bg "B_nav2" ros2 launch nav2_bringup navigation_launch.py use_sim_time:="$USE_SIM_TIME"
else
  start_bg "B_nav2" \
    ros2 launch nav2_bringup bringup_launch.py \
      use_sim_time:="$USE_SIM_TIME" \
      map:="$MAP_YAML" \
      slam:=false
fi

if [[ "$WITH_EXPLORE" == true && "$MODE" == "mapping" ]]; then
  start_bg "C_explore" ros2 launch explore_lite explore.launch.py use_sim_time:="$USE_SIM_TIME"
  last_index=$(( ${#PIDS[@]} - 1 ))
  EXPLORE_PID="${PIDS[$last_index]}"
  EXPLORE_LOG="${LOG_DIR}/C_explore.log"
fi

echo
echo "[run_all] started. Mode: ${MODE}"
echo "[run_all] logs: ${LOG_DIR}"
echo

if [[ "$MODE" == "mapping" ]]; then
  if [[ "$WITH_GZWEB" == true ]]; then
    if [[ "$NO_GZWEB_VITE" == true ]]; then
      echo "[run_all] UI: foundation /opt/gzweb on port 8000 (no office Vite)."
    elif wait_for_port 5173 25; then
      echo "[run_all] office gzweb-demo ready: http://localhost:5173"
    else
      echo "[run_all] WARN: office gzweb frontend not ready yet, check ${LOG_DIR}/E_gzweb_frontend.log"
    fi
  fi

  if [[ "${WITH_EXPLORE}" == true && -n "${EXPLORE_PID:-}" ]]; then
    echo "[run_all] exploration is running (pid=${EXPLORE_PID}), waiting for stop signal (timeout=${EXPLORE_WAIT_TIMEOUT}s)..."
    if wait_for_explore_stop_log_with_timer "${EXPLORE_LOG}" "${EXPLORE_WAIT_TIMEOUT}"; then
      echo "[run_all] detected exploration stop from log, finalizing explore process..."
    else
      echo "[run_all] WARN: explore stop not detected within timeout; forcing stop and saving current map."
    fi
    kill "${EXPLORE_PID}" 2>/dev/null || true
    wait "${EXPLORE_PID}" 2>/dev/null || true
    echo "[run_all] explore_lite stopped, now saving map."
  fi

  if [[ "${WITH_EXPLORE}" == true ]]; then
    echo "[run_all] auto-saving map to ${MAP_BASENAME}.yaml/.pgm ..."
    ros2 run nav2_map_server map_saver_cli -f "$MAP_BASENAME" --ros-args -p use_sim_time:="$USE_SIM_TIME"
    echo "[run_all] map saved successfully:"
    echo "  - ${MAP_BASENAME}.yaml"
    echo "  - ${MAP_BASENAME}.pgm"
    echo "[run_all] mapping workflow completed, exiting now."
    exit 0
  fi

  cat <<EOF
[run_all] Mapping mode notes:
  - RViz 'Localization: inactive' is normal during SLAM mapping.
  - You did not enable --explore, so auto-save/auto-exit is disabled.
  - Press Ctrl+C when you want to stop all.
EOF
  wait
else
  echo "[run_all] Localization mode: RViz Localization should become active after initial pose."
  echo "[run_all] Keep this terminal running; Ctrl+C to stop all."
  wait
fi