#!/usr/bin/env bash
# One GPU's share of the sweep. Co-hosts the OSCAR server + IsaacSim eval on the
# same GPU. Runs a list of "wm:task" jobs; (re)starts the server only when the WM
# changes (so order jobs by WM to get a single reload).
#
#   run_gpu_worker.sh <GPU> <OSCAR_PORT> <MASTER_PORT> <wm:task> [<wm:task> ...]
#     wm  = lora25 | base       task = open_drawer|open_laptop|pour_balls|push_box|stack_can|unload_cans
set -u
GPU=$1; OSCAR_PORT=$2; MASTER_PORT=$3; shift 3
JOBS=("$@")

HERE=/root/pwm_local/PWM_code/policy_wm_eval
OSCAR=/root/pwm_local/oscar-public
SERVER_PY=/root/pwm_local/PWM_code/policy_rl/rl_sim/act_wm_residual_oscar/workers/wm/oscar_wm_server.py
NUM_EPISODES=${NUM_EPISODES:-30}
OUT_DIR=${OUT_DIR:-$HERE/eval_out}
LOGDIR=${LOGDIR:-$OUT_DIR/logs}
mkdir -p "$LOGDIR"

lora_path() { case "$1" in lora25) echo "finetune/ckpt_6task_lora25/lora_step16000.pt";; base) echo "";; *) echo "??";; esac; }

SERVER_PGID=""
start_server() {   # $1 = wm
  local wm=$1 lora; lora=$(lora_path "$wm")
  echo "[gpu$GPU] start $wm server :$OSCAR_PORT (lora='${lora:-BASE}')" >&2
  setsid bash -c "cd '$OSCAR' && CUDA_VISIBLE_DEVICES=$GPU PYTHONPATH=. HF_HUB_OFFLINE=1 \
    .venv/bin/torchrun --nproc_per_node=1 --master_port=$MASTER_PORT \
    '$SERVER_PY' --lora '$lora' --lora_r 16 --lora_alpha 32 --port $OSCAR_PORT" \
    > "$LOGDIR/server_gpu${GPU}_${wm}.log" 2>&1 &
  SERVER_PGID=$!
  local i
  for i in $(seq 1 80); do
    if curl -s -m 3 "http://localhost:$OSCAR_PORT/health" 2>/dev/null | grep -q model_loaded; then
      echo "[gpu$GPU] $wm ready" >&2; return 0; fi
    kill -0 "$SERVER_PGID" 2>/dev/null || { echo "[gpu$GPU] server proc died" >&2; return 1; }
    sleep 15
  done
  echo "[gpu$GPU] server ready-timeout" >&2; return 1
}
stop_server() {
  [ -n "$SERVER_PGID" ] || return 0
  kill -TERM -"$SERVER_PGID" 2>/dev/null   # negative = whole process group (setsid leader)
  local i
  for i in $(seq 1 15); do kill -0 -"$SERVER_PGID" 2>/dev/null || break; sleep 2; done
  kill -KILL -"$SERVER_PGID" 2>/dev/null   # force any stragglers
  sleep 4                                   # let CUDA memory actually free before reloading
  SERVER_PGID=""
}
trap 'stop_server' EXIT

cur_wm=""
for job in "${JOBS[@]}"; do
  wm=${job%%:*}; task=${job##*:}
  wmtag=$([ "$wm" = lora25 ] && echo lora25_step16000 || echo base)
  if [ "$wm" != "$cur_wm" ]; then
    stop_server
    start_server "$wm" || { echo "[gpu$GPU] ABORT (server $wm failed)" >&2; exit 1; }
    cur_wm=$wm
  fi
  echo "[gpu$GPU] === eval $task in $wmtag ($NUM_EPISODES ep) ===" >&2
  TASK=$task OSCAR_PORT=$OSCAR_PORT WM_TAG=$wmtag GPU=$GPU \
    NUM_EPISODES=$NUM_EPISODES OUT_DIR=$OUT_DIR \
    bash "$HERE/run_eval.sh" > "$LOGDIR/eval_gpu${GPU}_${wmtag}_${task}.log" 2>&1
  echo "[gpu$GPU] finished $task/$wmtag rc=$?" >&2
done
echo "[gpu$GPU] ALL DONE" >&2