#!/usr/bin/env bash
# Start the slave Orin real-vehicle stack.

set -euo pipefail

GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

validate_bool() {
  local name="$1"
  local value="$2"

  case "${value}" in
    true|false)
      ;;
    *)
      log_error "${name} must be true or false, got: ${value}"
      exit 1
      ;;
  esac
}

script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

find_workspace_root() {
  if [[ -n "${SLAVE_WORKSPACE_ROOT:-}" ]]; then
    echo "${SLAVE_WORKSPACE_ROOT}"
    return
  fi

  local candidates=(
    "${script_dir}"
    "${PWD}"
  )
  for candidate in "${candidates[@]}"; do
    if [[ -f "${candidate}/install/setup.bash" ]]; then
      (cd "${candidate}" && pwd)
      return
    fi
  done

  echo "${PWD}"
}

find_vehicle_config_dir() {
  local vehicle_model="$1"
  local source_config="${WORKSPACE_ROOT}/src/Vehicle/${vehicle_model}_description/config/description"
  local install_config="${WORKSPACE_ROOT}/install/${vehicle_model}_description/share/${vehicle_model}_description/config/description"

  if [[ -d "${source_config}" ]]; then
    echo "${source_config}"
    return
  fi

  echo "${install_config}"
}

kill_process_group() {
  local pid="$1"
  local signal="${2:-TERM}"

  if [[ -z "${pid}" ]]; then
    return 0
  fi

  local pgid
  pgid="$(ps -o pgid= -p "${pid}" 2>/dev/null | tr -d ' ')"
  if [[ -z "${pgid}" ]]; then
    # Long-running commands are started with setsid, so the original pid is
    # also the process-group id even if the parent command exits first.
    pgid="${pid}"
  fi

  kill "-${signal}" -- "-${pgid}" 2>/dev/null || true
}

WORKSPACE_ROOT="$(find_workspace_root)"
DEFAULT_SLAVE_LAUNCH_FILE="${WORKSPACE_ROOT}/src/Launcher/farui_launch/launch/ecar_slave.launch.xml"
DEFAULT_HEARTBEAT_SCRIPT="${WORKSPACE_ROOT}/env_config/dual_orin/scripts/heartbeat_monitor.sh"
DEFAULT_VEHICLE_MODEL="ecar_vehicle"
ROS_SETUP_FILE="${SLAVE_ROS_SETUP_FILE:-/opt/ros/${ROS_DISTRO:-humble}/setup.bash}"

SLAVE_LAUNCH_FILE="${SLAVE_LAUNCH_FILE:-${DEFAULT_SLAVE_LAUNCH_FILE}}"
HEARTBEAT_SCRIPT="${SLAVE_HEARTBEAT_SCRIPT:-${DEFAULT_HEARTBEAT_SCRIPT}}"
VEHICLE_MODEL="${1:-${SLAVE_VEHICLE_MODEL:-${DEFAULT_VEHICLE_MODEL}}}"
CONFIG_DIR="${SLAVE_CONFIG_DIR:-$(find_vehicle_config_dir "${VEHICLE_MODEL}")}"
CAN_INTERFACE="${2:-${SLAVE_CAN_INTERFACE:-can1}}"
CAN_BITRATE="${3:-${SLAVE_CAN_BITRATE:-500000}}"
LAUNCH_HEARTBEAT="${SLAVE_LAUNCH_HEARTBEAT:-true}"
HEARTBEAT_DELAY="${SLAVE_HEARTBEAT_DELAY:-5}"
CONFIGURE_CAN="${SLAVE_CONFIGURE_CAN:-false}"
LAUNCH_SYSTEM="${SLAVE_LAUNCH_SYSTEM:-true}"
SYSTEM_RUN_MODE="${SLAVE_SYSTEM_RUN_MODE:-online}"
LAUNCH_SYSTEM_MONITOR="${SLAVE_LAUNCH_SYSTEM_MONITOR:-false}"
LAUNCH_DUMMY_DIAG_PUBLISHER="${SLAVE_LAUNCH_DUMMY_DIAG_PUBLISHER:-false}"

validate_bool "SLAVE_LAUNCH_HEARTBEAT" "${LAUNCH_HEARTBEAT}"
validate_bool "SLAVE_CONFIGURE_CAN" "${CONFIGURE_CAN}"
validate_bool "SLAVE_LAUNCH_SYSTEM" "${LAUNCH_SYSTEM}"
validate_bool "SLAVE_LAUNCH_SYSTEM_MONITOR" "${LAUNCH_SYSTEM_MONITOR}"
validate_bool "SLAVE_LAUNCH_DUMMY_DIAG_PUBLISHER" "${LAUNCH_DUMMY_DIAG_PUBLISHER}"

case "${SYSTEM_RUN_MODE}" in
  online|logging_simulation|planning_simulation)
    ;;
  *)
    log_error "SLAVE_SYSTEM_RUN_MODE must be online, logging_simulation, or planning_simulation; got: ${SYSTEM_RUN_MODE}"
    exit 1
    ;;
esac

if [[ ! -f "${WORKSPACE_ROOT}/install/setup.bash" ]]; then
  log_error "Workspace setup file not found: ${WORKSPACE_ROOT}/install/setup.bash"
  exit 1
fi

if [[ ! -f "${SLAVE_LAUNCH_FILE}" ]]; then
  log_error "Slave launch file not found: ${SLAVE_LAUNCH_FILE}"
  exit 1
fi

if [[ ! -f "${CONFIG_DIR}/sensors_calibration.yaml" || ! -f "${CONFIG_DIR}/sensor_kit_calibration.yaml" ]]; then
  log_error "Vehicle config files not found in: ${CONFIG_DIR}"
  log_error "Expected sensors_calibration.yaml and sensor_kit_calibration.yaml"
  exit 1
fi

if [[ "${LAUNCH_HEARTBEAT}" == "true" && ! -f "${HEARTBEAT_SCRIPT}" ]]; then
  log_error "Heartbeat script not found: ${HEARTBEAT_SCRIPT}"
  exit 1
fi

log_info "Workspace: ${WORKSPACE_ROOT}"
log_info "Slave launch: ${SLAVE_LAUNCH_FILE}"
log_info "Vehicle model: ${VEHICLE_MODEL}"
log_info "Vehicle config dir: ${CONFIG_DIR}"
log_info "CAN: ${CAN_INTERFACE}, bitrate: ${CAN_BITRATE}"
log_info "System: launch=${LAUNCH_SYSTEM}, run_mode=${SYSTEM_RUN_MODE}, monitor=${LAUNCH_SYSTEM_MONITOR}"

cd "${WORKSPACE_ROOT}"
set +u
if [[ -f "${ROS_SETUP_FILE}" ]]; then
  log_info "Sourcing ROS setup: ${ROS_SETUP_FILE}"
  source "${ROS_SETUP_FILE}"
else
  log_warn "ROS setup file not found, continuing with current shell environment: ${ROS_SETUP_FILE}"
fi
log_info "Sourcing workspace setup: ${WORKSPACE_ROOT}/install/setup.bash"
source "${WORKSPACE_ROOT}/install/setup.bash"
set -u

log_ros_pkg_prefix() {
  local package="$1"
  local prefix
  if prefix="$(ros2 pkg prefix "${package}" 2>/dev/null)"; then
    log_info "ROS package ${package}: ${prefix}"
  else
    log_warn "ROS package ${package} not found in the sourced environment."
  fi
}

log_ros_pkg_prefix "farui_launch"
log_ros_pkg_prefix "ecar_can_driver"
log_ros_pkg_prefix "autoware_trajectory_follower_node"
log_ros_pkg_prefix "system_diagnostic_monitor"
log_ros_pkg_prefix "component_state_monitor"

if [[ "${CONFIGURE_CAN}" == "true" ]]; then
  log_info "Configuring CAN interface ${CAN_INTERFACE}."
  sudo ip link set "${CAN_INTERFACE}" down 2>/dev/null || true
  sudo ip link set "${CAN_INTERFACE}" up type can bitrate "${CAN_BITRATE}"
else
  log_warn "CAN auto-configuration skipped. Set SLAVE_CONFIGURE_CAN=true to run sudo ip link setup."
fi

SLAVE_STACK_PID=""
HEARTBEAT_PID=""

log_info "Launching slave Orin stack."
setsid ros2 launch "${SLAVE_LAUNCH_FILE}" \
  vehicle_model:="${VEHICLE_MODEL}" \
  config_dir:="${CONFIG_DIR}" \
  can_interface:="${CAN_INTERFACE}" \
  can_bitrate:="${CAN_BITRATE}" \
  launch_system:="${LAUNCH_SYSTEM}" \
  system_run_mode:="${SYSTEM_RUN_MODE}" \
  launch_system_monitor:="${LAUNCH_SYSTEM_MONITOR}" \
  launch_dummy_diag_publisher:="${LAUNCH_DUMMY_DIAG_PUBLISHER}" &
SLAVE_STACK_PID=$!

if [[ "${LAUNCH_HEARTBEAT}" == "true" ]]; then
  log_info "Heartbeat monitor will launch in ${HEARTBEAT_DELAY} seconds."
  setsid bash -c '
    set -e
    delay="$1"
    heartbeat_script="$2"
    sleep "${delay}"
    echo -e "\033[0;32m[INFO]\033[0m Launching heartbeat monitor."
    exec "${heartbeat_script}"
  ' bash "${HEARTBEAT_DELAY}" "${HEARTBEAT_SCRIPT}" &
  HEARTBEAT_PID=$!
else
  log_info "Heartbeat monitor disabled by SLAVE_LAUNCH_HEARTBEAT=false."
fi

cleanup() {
  trap - EXIT INT TERM

  local pids=()
  if [[ -n "${HEARTBEAT_PID}" ]]; then
    pids+=("${HEARTBEAT_PID}")
  fi
  if [[ -n "${SLAVE_STACK_PID}" ]]; then
    pids+=("${SLAVE_STACK_PID}")
  fi

  if [[ ${#pids[@]} -eq 0 ]]; then
    return
  fi

  log_info "Shutting down launched processes."

  for pid in "${pids[@]}"; do
    kill_process_group "${pid}" TERM
  done

  sleep 2

  for pid in "${pids[@]}"; do
    kill_process_group "${pid}" KILL
  done

  wait "${pids[@]}" 2>/dev/null || true
}

trap cleanup EXIT INT TERM

WAIT_PIDS=()
if [[ -n "${SLAVE_STACK_PID}" ]]; then
  WAIT_PIDS+=("${SLAVE_STACK_PID}")
fi
if [[ -n "${HEARTBEAT_PID}" ]]; then
  WAIT_PIDS+=("${HEARTBEAT_PID}")
fi

wait "${WAIT_PIDS[@]}"