#!/usr/bin/env bash

set -euo pipefail

# Resolve workspace root (script may be invoked from repo root)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WS_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

# Source ROS 2 workspace (temporarily disable nounset to avoid unbound vars inside setup scripts)
if [[ -f "${WS_ROOT}/install/setup.bash" ]]; then
	# shellcheck source=/dev/null
	set +u
	source "${WS_ROOT}/install/setup.bash"
	set -u
else
	echo "[run.sh] install/setup.bash not found at ${WS_ROOT}/install/setup.bash"
	echo "[run.sh] Please build the workspace first, e.g.,: colcon build && source install/setup.bash"
	exit 1
fi

# Human-friendly logging format with local wall-clock time and colors
export RCUTILS_COLORIZED_OUTPUT=${RCUTILS_COLORIZED_OUTPUT:-1}
export RCUTILS_CONSOLE_OUTPUT_FORMAT='[{time}] [{severity}] [{name}]: {message}'

LOG_TO_FILES=0
LOG_LEVEL="info"   # debug|info|warn|error|fatal
RUN_BRAIN=1
RUN_MOTION=1

usage() {
	cat <<EOF
Usage: ${0##*/} [options]

Options:
	-l, --log-to-files        Write each node's stdout/stderr to log/run/<timestamp>/*.log
	-L, --log-level LEVEL     Set ROS 2 log level (default: info). Ex: debug, warn, error
			--brain-only          Run only brain node
			--arm_control-only         Run only arm_control node
	-h, --help                Show this help and exit

Environment:
	RCUTILS_COLORIZED_OUTPUT  Colorize console logs (default: 1)
	RCUTILS_CONSOLE_OUTPUT_FORMAT  Log format (already configured here)
EOF
}

# Parse args
while [[ $# -gt 0 ]]; do
	case "$1" in
		-l|--log-to-files)
			LOG_TO_FILES=1; shift ;;
		-L|--log-level)
			LOG_LEVEL="${2:-info}"; shift 2 ;;
		--brain-only)
			RUN_BRAIN=1; RUN_MOTION=0; shift ;;
		--arm_control-only)
			RUN_BRAIN=0; RUN_MOTION=1; shift ;;
		-h|--help)
			usage; exit 0 ;;
		*)
			echo "Unknown option: $1" >&2
			usage; exit 2 ;;
	esac
done

timestamp() { date +"%Y%m%d_%H%M%S"; }

LOG_DIR="${WS_ROOT}/log/run/$(timestamp)"
mkdir -p "${LOG_DIR}"

pids=()
names=()

start_node() {
	local name="$1"; shift
	local pkg="$1"; shift
	local exe="$1"; shift
	local log_file="${LOG_DIR}/${name}.log"

	if (( LOG_TO_FILES )); then
		echo "[run.sh] Starting ${name} -> ros2 run ${pkg} ${exe} (logs: ${log_file})"
		(
			# Disable ANSI colors in log files for readability
			RCUTILS_COLORIZED_OUTPUT=0 exec ros2 run "${pkg}" "${exe}" --ros-args --log-level "${LOG_LEVEL}"
		) >>"${log_file}" 2>&1 &
	else
		echo "[run.sh] Starting ${name} -> ros2 run ${pkg} ${exe} --ros-args --log-level ${LOG_LEVEL}"
		ros2 run "${pkg}" "${exe}" --ros-args --log-level "${LOG_LEVEL}" &
	fi
	local pid=$!
	pids+=("${pid}")
	names+=("${name}")
}

cleanup_done=0
cleanup() {
	(( cleanup_done )) && return 0
	cleanup_done=1
	echo ""
	echo "[run.sh] Shutting down (${#pids[@]} processes)..."
	# Send SIGINT first for graceful shutdown
	for pid in "${pids[@]:-}"; do
		if kill -0 "$pid" 2>/dev/null; then
			kill -INT "$pid" 2>/dev/null || true
		fi
	done
	# Give them a moment to exit gracefully
	sleep 1
	# Force kill any stubborn processes
	for pid in "${pids[@]:-}"; do
		if kill -0 "$pid" 2>/dev/null; then
			kill -TERM "$pid" 2>/dev/null || true
		fi
	done
	wait || true
	echo "[run.sh] All processes stopped."
}

trap cleanup INT TERM

# Start requested nodes
if (( RUN_MOTION )); then
	start_node "arm_control" "arm_control" "motion_node"
fi
if (( RUN_BRAIN )); then
	start_node "brain" "brain" "brain_node"
fi

if (( ${#pids[@]} == 0 )); then
	echo "[run.sh] Nothing to run. Use --brain-only or --arm_control-only to pick a node."
	exit 0
fi

echo "[run.sh] Started: ${names[*]}"
echo "[run.sh] Press Ctrl+C to stop."

# Wait for all background processes
wait

# Ensure cleanup executes on normal exit too
cleanup
