#!/usr/bin/env bash
# fleet_install.sh — deploy/manage a Sanad fleet agent on a robot over SSH,
# as FULL DOCKER with a user-level systemd auto-start service.
#
# INTERACTIVE (no args):   ./fleet_install.sh
#   → pick robot type → enter IP → it detects whether the agent is already
#     installed and offers the right actions (install / uninstall / data / ...).
#
# SCRIPTABLE — install requires: ip, --sn, --token, --server-url (and optionally --post):
#   ./fleet_install.sh install   <g1|r1|go2> <ip> --sn <robot-serial> --token <device-token> \
#                                --server-url https://fleet.example.com [--post /api/v1/fleet/ingest/telemetry] \
#                                [--name NAME] [--user U]
#
#   UPDATE an already-installed robot without re-entering the token — reuses the
#   DEVICE_TOKEN already in the robot's .env (it never leaves the robot):
#   ./fleet_install.sh install   <g1|r1|go2> <ip> --sn <robot-serial> --keep-token \
#                                --server-url https://fleet.example.com
#   ./fleet_install.sh uninstall <g1|r1|go2> <ip> [--user U]
#   ./fleet_install.sh status    <g1|r1|go2> <ip>
#   ./fleet_install.sh data      <g1|r1|go2> <ip>   # show the data it is sending
#   ./fleet_install.sh logs      <g1|r1|go2> <ip>
#   ./fleet_install.sh test      <g1|r1|go2> <ip> [--server-ip IP] [--keep-server]
#
# Deploy model: rsync the canonical workstation copy -> robot, `docker build`
# on the robot (native arm64), `docker create` the container, and a user
# systemd unit (linger-enabled) owns start/stop/auto-start. No docker-compose
# needed on the robot; no sudo needed (user-level systemd).
set -uo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
AGENTS="$SCRIPT_DIR/agents"
TYPES="g1 r1 go2"
declare -A KIND=( [g1]="TELEMETRY + MAP (unitree_hg)" [r1]="TELEMETRY (unitree_hg)" [go2]="TELEMETRY (unitree_go)" )

die(){ echo "ERROR: $*" >&2; exit 1; }
img_of(){ echo "sanad-api-$1"; }        # image + container + unit share this base
rdir_of(){ echo "sanad_api_$1"; }
unit_of(){ echo "sanad-api-$1.service"; }
sn_default(){ echo "${1}_${IP##*.}"; }

SSH_OPTS=(-o BatchMode=yes -o ConnectTimeout=8 -o StrictHostKeyChecking=accept-new)
rmt(){ ssh -n "${SSH_OPTS[@]}" "$USER_@$IP" "$@"; }        # -n: never read our stdin
rmt_in(){ ssh "${SSH_OPTS[@]}" "$USER_@$IP" "$@"; }        # stdin passthrough (heredocs)
detect_server_ip(){ [ -n "$SERVER_IP" ] && { echo "$SERVER_IP"; return; }
  ip -o route get "$IP" 2>/dev/null | grep -oP 'src \K\S+' | head -1; }

# ---- per-type docker create args ----
# Every agent now syncs maps: find where THIS robot's Sanad/SLAM stack keeps its
# maps on the host. Falls back to the agent's own dirs (agent then reports
# map: no_map until a map exists / is exposed).
MAPS_HOST=""; DATA_HOST=""; EXTRA_MAPS_HOST=""
probe_maps_dirs(){
  local t="$1"; local rdir; rdir="$(rdir_of "$t")"
  MAPS_HOST="/home/$USER_/$rdir/maps"
  DATA_HOST="/home/$USER_/$rdir/web_data"
  EXTRA_MAPS_HOST=""
  if [ "$t" = g1 ]; then
    for c in "/home/$USER_/marcus_nav2_test/maps" \
             "/home/$USER_/sanad_deploy/Sanad_Package_4/nav/data/ros"; do
      if rmt "test -d $c" 2>/dev/null; then MAPS_HOST="$c"; break; fi
    done
    for c in "/home/$USER_/marcus_nav2_test/web/data" \
             "/home/$USER_/sanad_deploy/Sanad_Package_4/nav/data/web"; do
      if rmt "test -d $c" 2>/dev/null; then DATA_HOST="$c"; break; fi
    done
  else
    # r1 (stereo VSLAM) / go2: find an rtabmap db or a map_server yaml set
    local found
    found=$(rmt "find /home/$USER_ -maxdepth 4 \( -name 'rtabmap*.db' -o -name '*.posegraph' \) 2>/dev/null | head -1" 2>/dev/null)
    [ -n "$found" ] && MAPS_HOST="$(dirname "$found")"
  fi
  # Nav2/Pudu deploy maps (map.pgm + map.yaml, incl. keepout-baked) usually live in
  # a *_nav2_docker/maps dir SEPARATE from the SLAM db above. Find it and mount it
  # as an extra map root so the agent discovers+uploads it too. Skip if it's the
  # same as the primary maps mount (g1 already maps marcus_nav2_test/maps).
  local e
  e=$(rmt "for d in /home/$USER_/${t}_nav2_docker/maps /home/$USER_/marcus_nav2_test/maps \
                     /home/$USER_/go2_nav2_docker/maps /home/$USER_/r1_nav2_docker/maps; do \
             ls \$d/*.yaml >/dev/null 2>&1 && { echo \$d; break; }; done" 2>/dev/null | head -1)
  if [ -n "$e" ] && [ "$e" != "$MAPS_HOST" ]; then EXTRA_MAPS_HOST="$e"; fi
}

run_args(){
  # /host:ro → real disk stats; maps/web_data ro + state rw → map sync
  local rdir; rdir="$(rdir_of "$1")"
  # extra Nav2/Pudu map dir (mounted at /data/nav2_maps, matches EXTRA_MAP_DIRS)
  local extra=""
  [ -n "$EXTRA_MAPS_HOST" ] && extra=" -v ${EXTRA_MAPS_HOST}:/data/nav2_maps:ro"
  echo "--network host --env-file /home/$USER_/$rdir/.env -v /:/host:ro \
    -v ${MAPS_HOST:-/home/$USER_/$rdir/maps}:/data/maps:ro \
    -v ${DATA_HOST:-/home/$USER_/$rdir/web_data}:/data/web_data:ro \
    -v /home/$USER_/$rdir/state:/data/state${extra}"
}

# Resolve the SERVER_URL + VERIFY_TLS for a deploy: a full --server-url (real
# HTTPS fleet server, TLS verified) wins; otherwise http://<detected-ip>:<port>
# (the local test server, TLS off).
server_url(){ [ -n "$SERVER_URL_OVERRIDE" ] && echo "$SERVER_URL_OVERRIDE" || echo "http://$1:$PORT"; }
verify_tls(){ [ -n "$VERIFY_TLS_OPT" ] && echo "$VERIFY_TLS_OPT" || { [ -n "$SERVER_URL_OVERRIDE" ] && echo 1 || echo 0; }; }

# ---- write the robot-side .env ----
push_env(){
  local t="$1" sip="$2" rdir surl vtls; rdir="$(rdir_of "$t")"
  surl="$(server_url "$sip")"; vtls="$(verify_tls)"
  # identity: brand fixed; type/model from the agent kind; friendly display
  # name defaults to <model>_<ip-last-octet> (e.g. r1_82, g1_58)
  local iface=eth0 btype=humanoid model="$t"
  [ "$t" = r1 ] && iface=eth10
  [ "$t" = go2 ] && btype=dog
  local rname="${NAME:-${model}_${IP##*.}}"
  # optional Sanad data dir (its size is shown on the dashboard) — probe
  # the known per-robot locations; path is as seen through /host (ro mount)
  local dpath=""
  for c in "/home/$USER_/SanadR1/data" "/home/$USER_/sanad_deploy/Sanad_Package_4/data"; do
    if rmt "test -d $c" 2>/dev/null; then dpath="/host$c"; break; fi
  done
  # --keep-token: updating an already-installed robot without re-entering the
  # device token. Snapshot the CURRENT .env on the robot; the old DEVICE_TOKEN
  # line is spliced back below. The token never leaves the robot.
  if [ "$KEEP_TOKEN" = 1 ]; then
    rmt "test -f ~/$rdir/.env && grep -q '^DEVICE_TOKEN=' ~/$rdir/.env" \
      || die "--keep-token: no existing ~/$rdir/.env with a DEVICE_TOKEN on $IP (pass --token instead)"
    rmt "cp ~/$rdir/.env ~/$rdir/.env.prev"
  fi
  # common telemetry env (all agents)
  rmt_in "cat > ~/$rdir/.env" <<EOF
SERVER_URL=$surl
DEVICE_TOKEN=$TOKEN
SN=$SN
ROBOT_NAME=$rname
ROBOT_BRAND=unitree
ROBOT_TYPE=$btype
ROBOT_MODEL=$model
STORAGE_DATA_PATH=$dpath
${POST_ENDPOINT:+TELEMETRY_ENDPOINT=$POST_ENDPOINT}
DDS_INTERFACE=$iface
DDS_DOMAIN=0
VERIFY_TLS=$vtls
POLL_INTERVAL=2
LOW_SOC=50
ROBOT=sanad
MAPS_DIR=/data/maps
EXTRA_MAP_DIRS=/data/nav2_maps
DATA_DIR=/data/web_data
STATE_DIR=/data/state
MAP_SELECT=all
MAP_UPLOAD_MODE=multipart
MAP_POLL_INTERVAL=30
LOGS_INTERVAL=60
SOFTWARE_ROS=foxy
PROJECT_LOG_CONTAINER=auto
REMOTE_ENABLE=1
CONTROL_ENABLE=0
ALERT_SCAN_INTERVAL=10
ALERT_LOG_COOLDOWN=300
EOF
  # splice the previous DEVICE_TOKEN back in — done entirely ON the robot, by
  # line surgery (no shell interpolation), so a token with any characters is
  # safe and it is never printed or transferred.
  if [ "$KEEP_TOKEN" = 1 ]; then
    rmt "cd ~/$rdir && grep -v '^DEVICE_TOKEN=' .env > .env.new \
         && grep '^DEVICE_TOKEN=' .env.prev >> .env.new \
         && mv .env.new .env && rm -f .env.prev" \
      || die "--keep-token: failed to preserve the existing token on $IP"
    say "   token:     kept the one already on the robot"
  fi
}

installed(){  # 0 = installed (unit file OR container present)
  rmt "test -f ~/.config/systemd/user/$(unit_of "$1") || docker container inspect $(img_of "$1") >/dev/null 2>&1" 2>/dev/null
}

do_install(){
  local t="$1" img rdir unit sip=""; img="$(img_of "$t")"; rdir="$(rdir_of "$t")"; unit="$(unit_of "$t")"
  # SN is the robot's REAL serial — it keys the robot on the fleet server, so it
  # must be entered at installation (no silent default).
  [ -n "$SN" ] || die "install requires the robot serial: --sn <robot-serial>  (e.g. --sn E39N4000Q6D7E70F)"
  if [ -z "$SERVER_URL_OVERRIDE" ]; then
    sip="$(detect_server_ip)"; [ -n "$sip" ] || die "cannot detect server IP toward $IP (use --server-ip or --server-url)"
  fi
  echo "== INSTALL $t ($SN) on $IP  ->  fleet server $(server_url "$sip") =="
  echo ">> rsync agent -> $USER_@$IP:~/$rdir/"
  rsync -az --delete --exclude '.env' --exclude '__pycache__' --exclude '*.pyc' \
    --exclude 'data/' --exclude 'state/' \
    -e "ssh ${SSH_OPTS[*]}" "$AGENTS/$t/" "$USER_@$IP:~/$rdir/" || die "rsync failed"
  rmt "mkdir -p ~/$rdir/maps/sanad ~/$rdir/web_data ~/$rdir/state ~/.config/systemd/user"
  push_env "$t" "$sip"
  probe_maps_dirs "$t"
  echo ">> maps dir:   $MAPS_HOST"
  echo ">> places dir: $DATA_HOST"
  echo ">> docker build on robot (native arm64; R1/Go2 compile DDS on first build) ..."
  rmt "cd ~/$rdir && docker build -t $img:latest ." || die "docker build failed"
  echo ">> creating container (systemd will own its lifecycle) ..."
  rmt "docker rm -f $img >/dev/null 2>&1; docker create --name $img $(run_args "$t") $img:latest >/dev/null"
  echo ">> installing user systemd service $unit (auto-start on boot via linger) ..."
  rmt_in "cat > ~/.config/systemd/user/$unit" <<EOF
[Unit]
Description=Sanad Fleet Agent ($t) -> $SN
After=docker.service network-online.target

[Service]
Type=simple
Restart=always
RestartSec=5
ExecStart=/usr/bin/docker start -a $img
ExecStop=/usr/bin/docker stop -t 10 $img

[Install]
WantedBy=default.target
EOF
  rmt "loginctl enable-linger \$USER >/dev/null 2>&1; \
       systemctl --user daemon-reload; \
       systemctl --user enable --now $unit"
  sleep 2
  echo ">> service state:"; rmt "systemctl --user --no-pager -l status $unit | sed -n '1,4p'"
  echo "== installed. show data: $0 data $t $IP =="
}

do_uninstall(){
  local t="$1" img unit rdir; img="$(img_of "$t")"; unit="$(unit_of "$t")"; rdir="$(rdir_of "$t")"
  echo "== UNINSTALL $t on $IP =="
  # reset-failed AFTER daemon-reload: stopping the unit leaves `docker start -a`
  # exiting non-zero, so systemd keeps a stale "not-found failed" entry in its
  # runtime state (visible in `systemctl --user --failed`) even once the unit
  # file is gone. Without this the uninstall looks like a broken service forever.
  rmt "systemctl --user disable --now $unit >/dev/null 2>&1; \
       rm -f ~/.config/systemd/user/$unit; systemctl --user daemon-reload; \
       systemctl --user reset-failed $unit >/dev/null 2>&1; \
       docker rm -f $img >/dev/null 2>&1; docker rmi $img:latest >/dev/null 2>&1; \
       rm -rf ~/$rdir; echo '   removed service, container, image, and ~/'$rdir"
  echo "== uninstalled =="
}

do_status(){
  local t="$1" img unit; img="$(img_of "$t")"; unit="$(unit_of "$t")"
  echo "== STATUS $t on $IP =="
  rmt "systemctl --user --no-pager status $unit 2>/dev/null | sed -n '1,4p' || echo 'no service'; echo; \
       docker ps -a --filter name=$img --format 'container: {{.Names}}  {{.Status}}  ({{.Image}})' || echo 'no container'"
}

do_logs(){ rmt "docker logs --tail 40 -f $(img_of "$1")"; }

do_data(){
  local t="$1" img; img="$(img_of "$t")"
  echo "== DATA $t on $IP  (what it is sending to the fleet server) =="
  rmt "systemctl --user --no-pager status $(unit_of "$t") 2>/dev/null | sed -n '1,3p'; echo '--- recent telemetry/map posts ---'; \
       docker logs --tail 12 $img 2>&1 | grep -E 'telemetry ok|uploaded|nothing to upload|heartbeat|POST failed' | tail -8 || docker logs --tail 12 $img 2>&1"
}

# ---- deploy-less end-to-end test against the workstation server ----
start_server(){ local reqlog="/tmp/fleet_reqs_${1}_$$.jsonl"; echo "$reqlog" > /tmp/fleet_last_reqlog
  PORT=$PORT REQLOG="$reqlog" python3 "$SCRIPT_DIR/fleet_test_server.py" >/tmp/fleet_srv_$$.log 2>&1 &
  echo $! > /tmp/fleet_srv_pid; sleep 1; }
stop_server(){ [ -f /tmp/fleet_srv_pid ] && kill "$(cat /tmp/fleet_srv_pid)" 2>/dev/null; rm -f /tmp/fleet_srv_pid; }

do_test(){
  local t="$1" img sip; img="$(img_of "$t")"
  sip="$(detect_server_ip)"; [ -n "$sip" ] || die "cannot detect server IP (use --server-ip)"
  echo "== TEST $t via $IP  (workstation server http://$sip:$PORT) =="
  rmt "docker image inspect $img:latest >/dev/null 2>&1" || { echo ">> not installed yet — installing first"; do_install "$t"; }
  push_env "$t" "$sip"
  start_server "$t"; trap 'stop_server' EXIT
  local rc; rc=$(rmt "curl -s -o /dev/null -w '%{http_code}' --max-time 6 http://$sip:$PORT/ping" || echo 000)
  [ "$rc" = 200 ] || die "robot cannot reach workstation server (http $rc) — firewall on $sip:$PORT?"
  echo "   reachability OK (HTTP $rc)"
  probe_maps_dirs "$t"
  local rdir; rdir="$(rdir_of "$t")"
  if [ "$MAPS_HOST" = "/home/$USER_/$rdir/maps" ]; then
    # no real maps dir on this robot — seed a fixture to prove the upload path
    rmt "mkdir -p ~/$rdir/maps/sanad ~/$rdir/web_data/sanad/places
         head -c 4096 /dev/urandom > ~/$rdir/maps/sanad/floor-test.db
         touch -d '10 minutes ago' ~/$rdir/maps/sanad/floor-test.db
         printf '%s' '{\"dock\":{\"x\":1.2,\"y\":3.4,\"qz\":0,\"qw\":1}}' > ~/$rdir/web_data/sanad/places/floor-test.json
         rm -f ~/$rdir/state/uploaded.json"
  fi
  echo ">> real state --once (--force map):"
  rmt "docker run --rm $(run_args "$t") $img:latest --once --force" || true
  echo ">> simulate --once:"
  rmt "docker run --rm $(run_args "$t") $img:latest --simulate --once" || true
  sleep 1; echo; echo ">> workstation server received:"
  local reqlog; reqlog="$(cat /tmp/fleet_last_reqlog)"; local result=0
  python3 - "$reqlog" "$t" <<'PY' && result=0 || result=$?
import json,sys,os
recs=[json.loads(l) for l in open(sys.argv[1])] if os.path.exists(sys.argv[1]) else []
t=sys.argv[2]
for r in recs:
    j=r.get("json")
    print("   POST",r["path"],"auth="+("yes" if r["auth"].startswith("Bearer") else "NO"))
    if j: print("       ",{k:j[k] for k in j if k!="db_base64"})
    if r.get("meta"): print("        meta:",r["meta"][:160])
# every agent now does telemetry + map; PASS needs both to have arrived
ok=(any(x["path"].endswith("/telemetry") for x in recs)
    and any(x["path"].endswith("/map") for x in recs))
print(f"\n   RESULT: {'PASS' if ok else 'FAIL'}  ({len(recs)} request(s))")
sys.exit(0 if ok else 1)
PY
  [ "${KEEP_SERVER:-0}" = 1 ] && { echo ">> --keep-server: server still up (pid $(cat /tmp/fleet_srv_pid 2>/dev/null))"; trap - EXIT; }
  return $result
}

# --------------------------- arg parsing --------------------------- #
CMD=""; ROBOT=""; IP=""; SERVER_IP=""; PORT=8799; TOKEN="test-token"; SN=""; NAME=""; USER_="unitree"; KEEP_SERVER=0; KEEP_TOKEN=0
SERVER_URL_OVERRIDE=""; VERIFY_TLS_OPT=""; POST_ENDPOINT=""
POSA=()
while [ $# -gt 0 ]; do case "$1" in
  --server-ip) SERVER_IP="$2"; shift 2;;
  --server-url) SERVER_URL_OVERRIDE="$2"; shift 2;;   # full URL (real HTTPS fleet server) → VERIFY_TLS=1
  --verify-tls) VERIFY_TLS_OPT="$2"; shift 2;;
  --port) PORT="$2"; shift 2;;
  --token) TOKEN="$2"; shift 2;;
  --keep-token) KEEP_TOKEN=1; shift;;         # update in place: reuse the token already on the robot
  --sn) SN="$2"; shift 2;;
  --name) NAME="$2"; shift 2;;              # friendly display name (default <model>_<last-octet>)
  --post) POST_ENDPOINT="$2"; shift 2;;     # ingest POST path (telemetry or map endpoint)
  --user) USER_="$2"; shift 2;;
  --keep-server) KEEP_SERVER=1; shift;;
  -h|--help) grep -E '^#( |$)' "$0" | sed 's/^# \{0,1\}//'; exit 0;;
  *) POSA+=("$1"); shift;;
esac; done
CMD="${POSA[0]:-}"; ROBOT="${POSA[1]:-}"; IP="${POSA[2]:-}"

# --------------------------- interactive --------------------------- #
choose_type(){
  echo "Which robot / agent?" >&2; local i=1; local -a arr=()
  for t in $TYPES; do echo "  $i) $t  — ${KIND[$t]}" >&2; arr[$i]="$t"; i=$((i+1)); done
  read -rp "Choice: " c
  if [[ "$c" =~ ^[0-9]+$ ]] && [ -n "${arr[$c]:-}" ]; then echo "${arr[$c]}"
  elif echo "$TYPES" | grep -qw "$c"; then echo "$c"
  else echo ""; fi
}

interactive(){
  echo "=== Sanad Fleet Agent Installer ==="
  ROBOT="$(choose_type)"; [ -n "$ROBOT" ] || die "invalid robot type"
  read -rp "Robot IP: " IP; [ -n "$IP" ] || die "no IP"
  read -rp "SSH user [$USER_]: " u; [ -n "$u" ] && USER_="$u"
  rmt true 2>/dev/null || die "cannot SSH to $USER_@$IP (key auth?)"
  if installed "$ROBOT"; then
    echo; echo ">> '$(img_of "$ROBOT")' is ALREADY INSTALLED on $IP."
    echo "  1) show data   2) status   3) logs   4) reinstall   5) UNINSTALL   6) quit"
    read -rp "Choice: " a
    case "$a" in
      1) do_data "$ROBOT";; 2) do_status "$ROBOT";; 3) do_logs "$ROBOT";;
      4) prompt_install; do_install "$ROBOT";;
      5) read -rp "Really uninstall $(img_of "$ROBOT") from $IP? [y/N] " y; [ "$y" = y ] && do_uninstall "$ROBOT" || echo "cancelled";;
      *) echo "bye";;
    esac
  else
    echo; echo ">> not installed yet — let's install."
    prompt_install
    do_install "$ROBOT"
  fi
}

prompt_install(){
  # SN = the robot's REAL serial; it keys the robot on the fleet server → required.
  while [ -z "$SN" ]; do
    read -rp "Robot serial SN (required, e.g. E39N4000PB89GF88): " SN
  done
  local defname="${ROBOT}_${IP##*.}"
  read -rp "Display name [$defname]: " nm; NAME="${nm:-$defname}"
  # Server: full URL (https://… → TLS verified) or a bare IP (local test server).
  local autos; autos="$(detect_server_ip)"
  read -rp "SERVER_URL (https://…) or server IP [${autos:-required}]: " si
  si="${si:-$autos}"
  case "$si" in
    http://*|https://*) SERVER_URL_OVERRIDE="$si";;
    "") die "server URL or IP required";;
    *)  SERVER_IP="$si"
        read -rp "Server port [$PORT]: " p; [ -n "$p" ] && PORT="$p";;
  esac
  read -rp "Device token [$TOKEN]: " tk; [ -n "$tk" ] && TOKEN="$tk"
  read -rp "POST endpoint [agent default]: " pe; [ -n "$pe" ] && POST_ENDPOINT="$pe"
}

# --------------------------- dispatch --------------------------- #
[ -z "$CMD" ] && { interactive; exit $?; }
[ -n "$ROBOT" ] && [ -n "$IP" ] || die "usage: $0 <install|uninstall|status|data|logs|test> <g1|r1|go2> <ip> [opts]  (or run with no args for interactive)"
echo "$TYPES" | grep -qw "$ROBOT" || die "robot must be one of: $TYPES"
[ -d "$AGENTS/$ROBOT" ] || die "agent dir missing: $AGENTS/$ROBOT"
# SN: REQUIRED for install (the robot's real serial — checked in do_install);
# for test-only runs a placeholder default is fine.
if [ -z "$SN" ] && [ "$CMD" != install ]; then SN="$(sn_default "$ROBOT")"; fi

case "$CMD" in
  install)   do_install "$ROBOT";;
  uninstall) do_uninstall "$ROBOT";;
  status)    do_status "$ROBOT";;
  data)      do_data "$ROBOT";;
  logs)      do_logs "$ROBOT";;
  test)      do_test "$ROBOT";;
  *) die "unknown command: $CMD (install|uninstall|status|data|logs|test)";;
esac