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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WS_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
STATE_DIR="/tmp/lbot_arm_control"
PID_FILE="${STATE_DIR}/lbot_driver.pid"
PGID_FILE="${STATE_DIR}/lbot_driver.pgid"
BAG_PID_FILE="${STATE_DIR}/rosbag.pid"
BAG_PGID_FILE="${STATE_DIR}/rosbag.pgid"
LOG_FILE="${WS_DIR}/log/lbot_arm_control.log"
BAG_LOG_FILE="${WS_DIR}/log/lbot_arm_rosbag.log"

usage() {
  echo "Usage: $0 {start|stop|status}"
}

is_running() {
  [[ -f "${PID_FILE}" ]] || return 1
  local pid
  pid="$(cat "${PID_FILE}")"
  [[ -n "${pid}" ]] || return 1
  kill -0 "${pid}" 2>/dev/null
}

is_bag_running() {
  [[ -f "${BAG_PID_FILE}" ]] || return 1
  local pid
  pid="$(cat "${BAG_PID_FILE}")"
  [[ -n "${pid}" ]] || return 1
  kill -0 "${pid}" 2>/dev/null
}

source_workspace() {
  set +u
  if [[ -n "${ROS_DISTRO:-}" && -f "/opt/ros/${ROS_DISTRO}/setup.bash" ]]; then
    # shellcheck disable=SC1090
    source "/opt/ros/${ROS_DISTRO}/setup.bash"
  elif [[ -f "/opt/ros/jazzy/setup.bash" ]]; then
    # shellcheck disable=SC1091
    source "/opt/ros/jazzy/setup.bash"
  fi
  set -u

  if [[ ! -f "${WS_DIR}/install/setup.bash" ]]; then
    echo "Workspace is not built: ${WS_DIR}/install/setup.bash not found"
    echo "Run: cd ${WS_DIR} && colcon build --packages-select lbot_arm_interfaces lbot_driver"
    exit 1
  fi

  # shellcheck disable=SC1091
  set +u
  source "${WS_DIR}/install/setup.bash"
  set -u
}

start_bag_recorder() {
  mkdir -p "${STATE_DIR}" "$(dirname "${BAG_LOG_FILE}")"

  if is_bag_running; then
    echo "rosbag recorder is already running, pid=$(cat "${BAG_PID_FILE}")"
    return 0
  fi

  echo "Starting rosbag recorder..."
  echo "Rosbag log: ${BAG_LOG_FILE}"

  : > "${BAG_LOG_FILE}"
  setsid env WS_DIR="${WS_DIR}" bash -lc '
    set -euo pipefail
    set +u
    source "${WS_DIR}/install/setup.bash"
    set -u
    while true; do
      start_time=$(date +%Y%m%d_%H%M%S)
      output=/tmp/lbot_arm_${start_time}.mcap
      echo "Starting rosbag: ${output}"
      set +e
      timeout 3600 ros2 bag record \
        --storage mcap \
        --output "${output}" \
        --polling-interval 2 \
        --include-unpublished-topics \
        --disable-keyboard-controls \
        --log-level warn \
        --topics \
          /robot1/left_arm/joint_states \
          /robot1/right_arm/joint_states \
          /robot1/system_error \
          /robot1/left_hand/set_l20_joint \
          /robot1/right_hand/set_l20_joint \
        --services \
          /robot1/left_arm/move_joint \
          /robot1/right_arm/move_joint
      rc=$?
      set -e
      if [[ ${rc} -ne 124 ]]; then
        exit ${rc}
      fi
    done
  ' >> "${BAG_LOG_FILE}" 2>&1 &

  local pid=$!
  local pgid
  pgid="$(ps -o pgid= -p "${pid}" | tr -d ' ')"
  echo "${pid}" > "${BAG_PID_FILE}"
  echo "${pgid:-${pid}}" > "${BAG_PGID_FILE}"

  sleep 1
  if is_bag_running; then
    echo "rosbag recorder started, pid=${pid}, pgid=$(cat "${BAG_PGID_FILE}")"
  else
    echo "rosbag recorder failed to start. Check log: ${BAG_LOG_FILE}"
    rm -f "${BAG_PID_FILE}" "${BAG_PGID_FILE}"
    return 1
  fi
}

stop_bag_recorder() {
  if ! [[ -f "${BAG_PID_FILE}" ]]; then
    return 0
  fi

  local pid pgid
  pid="$(cat "${BAG_PID_FILE}")"
  pgid="${pid}"
  if [[ -f "${BAG_PGID_FILE}" ]]; then
    pgid="$(cat "${BAG_PGID_FILE}")"
  fi

  if ! kill -0 "${pid}" 2>/dev/null; then
    rm -f "${BAG_PID_FILE}" "${BAG_PGID_FILE}"
    return 0
  fi

  echo "Stopping rosbag recorder, pid=${pid}, pgid=${pgid}..."
  kill -TERM -- "-${pgid}" 2>/dev/null || kill -TERM "${pid}" 2>/dev/null || true

  for _ in {1..30}; do
    if ! kill -0 "${pid}" 2>/dev/null; then
      rm -f "${BAG_PID_FILE}" "${BAG_PGID_FILE}"
      echo "rosbag recorder stopped"
      return 0
    fi
    sleep 0.2
  done

  echo "rosbag recorder did not exit after SIGTERM, forcing stop..."
  kill -KILL -- "-${pgid}" 2>/dev/null || kill -KILL "${pid}" 2>/dev/null || true
  rm -f "${BAG_PID_FILE}" "${BAG_PGID_FILE}"
  echo "rosbag recorder stopped"
}

start_driver() {
  mkdir -p "${STATE_DIR}" "$(dirname "${LOG_FILE}")"

  if is_running; then
    echo "lbot_driver is already running, pid=$(cat "${PID_FILE}")"
    start_bag_recorder
    return 0
  fi

  source_workspace

  echo "Starting lbot_driver from officer_sdk..."
  echo "Log: ${LOG_FILE}"

  : > "${LOG_FILE}"
  setsid bash -lc "set +u; source '${WS_DIR}/install/setup.bash'; set -u; exec ros2 launch lbot_driver lbot_start_driver.launch.py" \
    >> "${LOG_FILE}" 2>&1 &

  local pid=$!
  local pgid
  pgid="$(ps -o pgid= -p "${pid}" | tr -d ' ')"
  echo "${pid}" > "${PID_FILE}"
  echo "${pgid:-${pid}}" > "${PGID_FILE}"

  sleep 1
  if is_running; then
    echo "lbot_driver started, pid=${pid}, pgid=$(cat "${PGID_FILE}")"
    start_bag_recorder
  else
    echo "lbot_driver failed to start. Check log: ${LOG_FILE}"
    rm -f "${PID_FILE}" "${PGID_FILE}"
    exit 1
  fi
}

stop_driver() {
  if ! [[ -f "${PID_FILE}" ]]; then
    echo "lbot_driver is not running"
    stop_bag_recorder
    return 0
  fi

  local pid pgid
  pid="$(cat "${PID_FILE}")"
  pgid="${pid}"
  if [[ -f "${PGID_FILE}" ]]; then
    pgid="$(cat "${PGID_FILE}")"
  fi

  if ! kill -0 "${pid}" 2>/dev/null; then
    echo "lbot_driver process is not alive, cleaning state"
    rm -f "${PID_FILE}" "${PGID_FILE}"
    stop_bag_recorder
    return 0
  fi

  echo "Stopping lbot_driver, pid=${pid}, pgid=${pgid}..."
  kill -TERM -- "-${pgid}" 2>/dev/null || kill -TERM "${pid}" 2>/dev/null || true

  for _ in {1..30}; do
    if ! kill -0 "${pid}" 2>/dev/null; then
      rm -f "${PID_FILE}" "${PGID_FILE}"
      echo "lbot_driver stopped"
      stop_bag_recorder
      return 0
    fi
    sleep 0.2
  done

  echo "lbot_driver did not exit after SIGTERM, forcing stop..."
  kill -KILL -- "-${pgid}" 2>/dev/null || kill -KILL "${pid}" 2>/dev/null || true
  rm -f "${PID_FILE}" "${PGID_FILE}"
  echo "lbot_driver stopped"
  stop_bag_recorder
}

status_driver() {
  if is_running; then
    echo "lbot_driver is running, pid=$(cat "${PID_FILE}"), pgid=$(cat "${PGID_FILE}" 2>/dev/null || cat "${PID_FILE}")"
  else
    echo "lbot_driver is not running"
  fi

  if is_bag_running; then
    echo "rosbag recorder is running, pid=$(cat "${BAG_PID_FILE}"), pgid=$(cat "${BAG_PGID_FILE}" 2>/dev/null || cat "${BAG_PID_FILE}")"
  else
    echo "rosbag recorder is not running"
  fi
}

case "${1:-}" in
  start)
    start_driver
    ;;
  stop)
    stop_driver
    ;;
  status)
    status_driver
    ;;
  *)
    usage
    exit 1
    ;;
esac