#!/bin/bash
# Per-task driver. Usage: run_tasks.sh <model-key> <outdir> [reps]
# For each rep: every config gets a fresh server (fresh temp HOME, so no prefix/response
# cache can carry over between reps), a discarded warm-up, then each task once.
# Config order rotates every rep (A B C, B C A, C A B). Only PIDs started here are signalled.
set -u
W=$HOME/rmlx-vsmlx; source $W/lib.sh; mkdir -p $W/homes
KEY=$1; OUT=$2; REPS=${3:-3}; ONLY=${ONLY:-}
case $KEY in
  4b)  ALIAS=qwen3.5-4b-4bit; SN=Qwen3.5-4B-MLX-4bit ;;
  9b)  ALIAS=qwen3.5-9b-4bit; SN=Qwen3.5-9B-4bit ;;
  35b) ALIAS=qwen3.6-35b-4bit; SN=Qwen3.6-35B-A3B-4bit ;;
esac
# CONFIGS: space-separated names; extra args per name below
CONFIGS=${CONFIGS:-"mlxlm rapid"}
extra_for() {
  case $1 in
    rapid_mtp)  echo '--speculative-config {"method":"mtp"}' ;;
    rapid_mtpk3) echo '--speculative-config {"method":"mtp","num_speculative_tokens":3}' ;;
    rapid_nospec) echo '--no-spec-decode' ;;
    *) echo '' ;;
  esac
}
mkdir -p $OUT
read -r -a CL <<<"$CONFIGS"; N=${#CL[@]}
{ date; sw_vers -productVersion; sysctl -n machdep.cpu.brand_string hw.memsize; echo "model=$KEY alias=$ALIAS snap=$(snap $SN)"; echo "configs=$CONFIGS reps=$REPS only=$ONLY"; } >>$OUT/env.txt
for rep in $(seq 1 $REPS); do
  for j in $(seq 0 $((N-1))); do
    cfg=${CL[$(( (j + rep - 1) % N ))]}
    eng=rapid; [ $cfg = mlxlm ] && eng=mlxlm
    cooldown
    busy_check || exit 1
    echo "$(date '+%F %T') rep=$rep cfg=$cfg load=$(sysctl -n vm.loadavg)" | tee -a $OUT/progress.txt
    # shellcheck disable=SC2046
    if [ $eng = rapid ] && [ -n "$(extra_for $cfg)" ]; then
      X=$(extra_for $cfg); set -f; read -r -a XA <<<"$X"; set +f
      start_server rapid $ALIAS $SN 18921 $OUT/$cfg "${XA[@]}" || { echo "start fail $cfg" | tee -a $OUT/progress.txt; stop_server; continue; }
    else
      start_server $eng $ALIAS $SN 18921 $OUT/$cfg || { echo "start fail $cfg" | tee -a $OUT/progress.txt; stop_server; continue; }
    fi
    peak_rss_sampler $PID $OUT/$cfg/rss-rep$rep.txt
    $W/venv-mlxlm/bin/python $W/bench_tasks.py --base-url $BASE --model "$MODEL" --label $cfg --tasks $W/tasks.json --out $OUT/$cfg/rep$rep.json ${ONLY:+--only $ONLY} >>$OUT/$cfg/client.log 2>&1
    stop_server
  done
done
echo "$(date '+%F %T') DONE $KEY" | tee -a $OUT/progress.txt