#!/usr/bin/env bash
# start_robot_server.sh — Set up SSH tunnel + launch ROS2 services for Panda joint state streaming.
#
# Usage:
#   ./scripts/start_robot_server.sh <[user@]ngrok_host> <ngrok_port>
#
# Examples:
#   ./scripts/start_robot_server.sh 4.tcp.us-cal-1.ngrok.io 12345
#   ./scripts/start_robot_server.sh root@4.tcp.us-cal-1.ngrok.io 12345
#
# Safe to re-run — kills existing sessions and remote processes before starting fresh.
# Starts 3 tmux sessions (stay open on failure so you can see errors):
#   panda_tunnel     — local SSH tunnel forwarding localhost:9090 -> remote:9090
#   panda_rosbridge  — rosbridge_websocket on the remote machine
#   panda_driver     — Franka FR3 driver on the remote machine

ARG_HOST="${1:?Usage: $0 <[user@]ngrok_host> <ngrok_port>}"
NGROK_PORT="${2:?Usage: $0 <[user@]ngrok_host> <ngrok_port>}"

# Handle both "host" and "user@host" formats
if [[ "$ARG_HOST" == *"@"* ]]; then
    REMOTE_USER="${ARG_HOST%%@*}"
    NGROK_HOST="${ARG_HOST##*@}"
else
    REMOTE_USER="root"
    NGROK_HOST="$ARG_HOST"
fi

REMOTE_PASS="elitesmasher99"
SSH_OPTS="-o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/dev/null -o ConnectTimeout=10"
SSH_BASE="sshpass -p '${REMOTE_PASS}' ssh ${SSH_OPTS} ${REMOTE_USER}@${NGROK_HOST} -p ${NGROK_PORT}"

echo "==> Connecting to ${REMOTE_USER}@${NGROK_HOST}:${NGROK_PORT}"

# Kill local tmux sessions so re-runs start clean
echo "==> Killing any existing panda tmux sessions..."
tmux kill-session -t panda_tunnel    2>/dev/null || true
tmux kill-session -t panda_rosbridge 2>/dev/null || true
tmux kill-session -t panda_driver    2>/dev/null || true

# Kill stale remote processes via a dedicated cleanup SSH call.
# Uses "python3.*rosbridge" pattern (matches the actual python process, not this script).
# SIGKILL (-9) ensures they die immediately. || true so we continue if nothing was running.
echo "==> Cleaning up stale remote processes..."
sshpass -p "${REMOTE_PASS}" ssh ${SSH_OPTS} ${REMOTE_USER}@${NGROK_HOST} -p ${NGROK_PORT} \
    'pkill -9 -f "python3.*rosbridge" 2>/dev/null; pkill -9 -f "franka.launch.py" 2>/dev/null; sleep 1' || true
sleep 1  # Give ports a moment to be freed

echo "==> [1/3] Starting SSH tunnel (localhost:9090 -> remote:9090)..."
tmux new-session -d -s panda_tunnel -x 220 -y 50 \
    "sshpass -p '${REMOTE_PASS}' ssh ${SSH_OPTS} -N -L 9090:localhost:9090 ${REMOTE_USER}@${NGROK_HOST} -p ${NGROK_PORT}; echo '--- tunnel exited (code $?) ---'; read"

echo "==> [2/3] Starting rosbridge on remote..."
tmux new-session -d -s panda_rosbridge -x 220 -y 50 \
    "${SSH_BASE} 'source /opt/ros/humble/setup.bash && ros2 launch rosbridge_server rosbridge_websocket_launch.xml'; echo '--- rosbridge exited (code $?) ---'; read"

echo "==> [3/3] Starting Franka FR3 driver on remote..."
tmux new-session -d -s panda_driver -x 220 -y 50 \
    "${SSH_BASE} 'source /opt/ros/humble/setup.bash && source /root/franka_ws/install/setup.bash && ros2 launch franka_bringup franka.launch.py robot_ip:=169.254.2.83 arm_id:=fr3 load_gripper:=true'; echo '--- franka driver exited (code $?) ---'; read"

echo ""
echo "==> All sessions launched. Attach to inspect output:"
echo "    tmux attach -t panda_tunnel"
echo "    tmux attach -t panda_rosbridge"
echo "    tmux attach -t panda_driver"
echo ""
echo "    Wait for rosbridge + Franka driver to finish starting, then run:"
echo "    python stream_panda_with_vis.py"