#!/usr/bin/env bash
# ================================================================
# ~20h benchmark matching dsv4_pro6000_vllm_tp_dp_matrix matrix.json
# Per-ISL TTFT SLO: 7s(min) -> 40s(max), stops concurrency sweep on violation
# OOM fault tolerance: auto-detect, skip combo, auto-restart server
# 3 repetitions for statistical robustness
# ================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/config.env"

RUN_ID="${RUN_ID:-$(date +%Y%m%d_%H%M%S)}"
RESULT_DIR="${RESULT_BASE}/${RUN_ID}"
SUMMARY_FILE="${RESULT_DIR}/summary.csv"
mkdir -p "$RESULT_DIR"
echo "isl,osl,concurrency,tps,ttft_ms,tpot_ms,e2e_ms,slo_ms,status,elapsed_s" >> "$SUMMARY_FILE"

log() { echo "[$(date '+%H:%M:%S')] $*"; }

# ================================================================
# TTFT SLO per ISL (DeepSeek-V4-Pro, TP=16 cross-node, min=7000ms)
# ================================================================
get_slo_ms() {
  local isl="$1"
  if   (( isl <= 2048 ));   then echo 7000
  elif (( isl <= 8192 ));   then echo 10000
  elif (( isl <= 32768 ));  then echo 15000
  elif (( isl <= 131072 )); then echo 25000
  else                           echo 40000
  fi
}

# ================================================================
# Reference matrix from dsv4_pro6000_vllm_tp_dp_matrix/matrix.json
# ================================================================
declare -A MATRIX
MATRIX["1024,128"]=Y;    MATRIX["1024,256"]=Y;    MATRIX["1024,512"]=Y
MATRIX["1024,1024"]=Y;   MATRIX["1024,2048"]=Y;   MATRIX["1024,4096"]=Y
MATRIX["4096,128"]=Y;    MATRIX["4096,256"]=Y;    MATRIX["4096,512"]=Y
MATRIX["4096,1024"]=Y;   MATRIX["4096,2048"]=Y;   MATRIX["4096,4096"]=Y
MATRIX["8192,128"]=Y;    MATRIX["8192,256"]=Y;    MATRIX["8192,512"]=Y
MATRIX["8192,1024"]=Y;   MATRIX["8192,2048"]=Y;   MATRIX["8192,4096"]=Y
MATRIX["16384,128"]=Y;   MATRIX["16384,256"]=Y;   MATRIX["16384,512"]=Y
MATRIX["16384,1024"]=Y;  MATRIX["16384,2048"]=Y;  MATRIX["16384,4096"]=P
MATRIX["32768,128"]=Y;   MATRIX["32768,256"]=Y;   MATRIX["32768,512"]=Y
MATRIX["32768,1024"]=Y;  MATRIX["32768,2048"]=Y;  MATRIX["32768,4096"]=P
MATRIX["65536,128"]=Y;   MATRIX["65536,256"]=Y;   MATRIX["65536,512"]=Y
MATRIX["65536,1024"]=Y;  MATRIX["65536,2048"]=P;  MATRIX["65536,4096"]=N
MATRIX["131072,128"]=Y;  MATRIX["131072,256"]=Y;  MATRIX["131072,512"]=Y
MATRIX["131072,1024"]=P; MATRIX["131072,2048"]=N; MATRIX["131072,4096"]=N
MATRIX["262144,128"]=Y;  MATRIX["262144,256"]=Y;  MATRIX["262144,512"]=P
MATRIX["262144,1024"]=N; MATRIX["262144,2048"]=N; MATRIX["262144,4096"]=N
MATRIX["524288,128"]=Y;  MATRIX["524288,256"]=P;  MATRIX["524288,512"]=N
MATRIX["524288,1024"]=N; MATRIX["524288,2048"]=N; MATRIX["524288,4096"]=N
MATRIX["1048576,128"]=Y; MATRIX["1048576,256"]=P; MATRIX["1048576,512"]=N
MATRIX["1048576,1024"]=N; MATRIX["1048576,2048"]=N; MATRIX["1048576,4096"]=N

ISL_LIST=(1024 4096 8192 16384 32768 65536 131072 262144 524288 1048576)
OSL_LIST=(128 256 512 1024 2048 4096)

get_conc_for() {
  local isl="$1"
  if   (( isl <= 4096 ));   then echo "1 2 4 8 16 32 64 128"
  elif (( isl <= 16384 ));  then echo "1 2 4 8 16 32"
  elif (( isl <= 65536 ));  then echo "1 2 4 8"
  else                       echo "1 2 4"
  fi
}

# ================================================================
# Benchmark runner with per-ISL SLO + OOM tolerance
# ================================================================
run_one() {
  local isl="$1" osl="$2" conc="$3"
  local slo_ms; slo_ms=$(get_slo_ms "$isl")
  local num_prompts=$(( conc * 5 ))
  [[ $num_prompts -lt 10 ]] && num_prompts=10
  local out="${RESULT_DIR}/isl${isl}_osl${osl}_c${conc}.json"
  local logf="${RESULT_DIR}/isl${isl}_osl${osl}_c${conc}.log"
  local start_ts=$(date +%s)

  docker run --rm --network host \
    -v "${MODEL_PATH}:${MODEL_PATH}:ro" \
    -v "${DATASET_PATH}:${DATASET_PATH}:ro" \
    -v "${RESULT_DIR}:${RESULT_DIR}" \
    -e PYTHONUNBUFFERED=1 -e HF_HUB_OFFLINE=1 -e TRANSFORMERS_OFFLINE=1 \
    --entrypoint python3 "$DOCKER_CLIENT_IMAGE" \
    -m "${SGLANG_BENCH_MODULE}" \
    --backend sglang --host "${HEAD_IP}" --port "${SGLANG_PORT}" \
    --dataset-name "${BENCH_DATASET_NAME}" --dataset-path "${DATASET_PATH}" \
    --random-input-len "$isl" --random-output-len "$osl" \
    --random-range-ratio "${RANDOM_RANGE_RATIO}" \
    --num-prompts "$num_prompts" --max-concurrency "$conc" \
    --request-rate 10000 --warmup-requests "${WARMUP_REQUESTS}" \
    --output-file "$out" --output-details --disable-tqdm \
    > "$logf" 2>&1
  local rc=$?
  local elapsed=$(($(date +%s) - start_ts))

  # OOM detection
  if [[ $rc -ne 0 ]]; then
    if grep -qi "out of memory\|CUDA error\|RESOURCE_EXHAUSTED" "$logf" 2>/dev/null; then
      echo "${isl},${osl},${conc},-1,-1,-1,-1,${slo_ms},OOM,${elapsed}" >> "$SUMMARY_FILE"
      log "  OOM  isl=${isl} osl=${osl} c=${conc}"; return 2
    fi
    echo "${isl},${osl},${conc},-1,-1,-1,-1,${slo_ms},FAIL,${elapsed}" >> "$SUMMARY_FILE"
    log "  FAIL isl=${isl} osl=${osl} c=${conc} rc=${rc}"; return 1
  fi

  # Extract metrics (JSONL: first line has aggregate summary)
  local metrics=$(python3 -c "
import json
for line in open('${out}'):
    line=line.strip()
    if not line: continue
    try:
        d=json.loads(line)
        tps=d.get('total_throughput',d.get('tps',-1))
        ttft=d.get('mean_ttft_ms',d.get('ttft_mean',-1))
        tpot=d.get('mean_tpot_ms',d.get('tpot_mean',-1))
        e2e=d.get('mean_e2e_latency_ms',d.get('e2e_latency_mean',-1))
        if tps!=-1 or ttft!=-1:
            print(f'{round(tps,1)},{round(ttft,1)},{round(tpot,1)},{round(e2e,1)}')
            break
    except: pass
" 2>/dev/null || echo "-1,-1,-1,-1")
  local tps=$(echo "$metrics" | cut -d, -f1)
  local ttft=$(echo "$metrics" | cut -d, -f2)
  local tpot=$(echo "$metrics" | cut -d, -f3)
  local e2e=$(echo "$metrics" | cut -d, -f4)

  echo "${isl},${osl},${conc},${tps},${ttft},${tpot},${e2e},${slo_ms},OK,${elapsed}" >> "$SUMMARY_FILE"

  # TTFT SLO check (per-ISL)
  if [[ "$ttft" != "-1" ]]; then
    if (( $(echo "$ttft > ${slo_ms}" | bc -l 2>/dev/null) )); then
      log "  SLO! ttft=${ttft}ms > slo=${slo_ms}ms -> stop higher conc"
      return 3
    fi
  fi

  log "  OK   isl=${isl} osl=${osl} c=${conc} tps=${tps} ttft=${ttft}ms slo=${slo_ms}ms"
  return 0
}

# ================================================================
# OOM recovery: restart server if it died
# ================================================================
check_and_recover_server() {
  if curl --fail --silent --max-time 5 "http://${HEAD_IP}:${SGLANG_PORT}/health" >/dev/null 2>&1; then
    return 0
  fi
  log "WARN: Server health failed. Checking..."
  local hok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${HEAD_NODE}" \
    "docker ps --filter name=${EXPERIMENT}_head --format '{{.Names}}' 2>/dev/null" || echo "")
  local wok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${WORKER_NODE}" \
    "docker ps --filter name=${EXPERIMENT}_worker --format '{{.Names}}' 2>/dev/null" || echo "")
  if [[ -z "$hok" || -z "$wok" ]]; then
    log "ERROR: Containers died. Restarting server..."
    bash "${SCRIPT_DIR}/start_sglang_multinode.sh" || { log "FATAL: Restart failed"; return 1; }
    log "Server restarted."
  fi
  return 0
}

# ================================================================
# Main
# ================================================================
main() {
  log "=== ${EXPERIMENT} Benchmark ==="
  log "Matrix: dsv4_pro6000_vllm_tp_dp_matrix"
  log "TTFT SLO: per-ISL (7s-40s), Reps: ${NUM_REPETITIONS}"
  log "Server: ${HEAD_IP}:${SGLANG_PORT}, Results: ${RESULT_DIR}"

  if ! curl --fail --silent --max-time 5 "http://${HEAD_IP}:${SGLANG_PORT}/health" >/dev/null 2>&1; then
    log "FATAL: Server not healthy"; exit 1
  fi
  log "Server health OK."

  local total=0 ok=0 fail=0 oom=0 slo=0 skip=0
  local start_time=$(date +%s)

  for rep in $(seq 1 ${NUM_REPETITIONS}); do
  log "========== REPETITION ${rep}/${NUM_REPETITIONS} =========="

  for isl in "${ISL_LIST[@]}"; do
    if ! check_and_recover_server; then exit 1; fi
    for osl in "${OSL_LIST[@]}"; do
      local tag="${MATRIX[${isl},${osl}]:-N}"
      [[ "$tag" == "N" ]] && { ((skip++)); continue; }
      log "--- isl=${isl} osl=${osl} [${tag}] ---"

      local conc_list=($(get_conc_for "$isl"))
      local combo_oom=0 combo_slo=0

      for conc in "${conc_list[@]}"; do
        [[ $combo_oom -eq 1 ]] && { ((skip++)); continue; }
        [[ $combo_slo -eq 1 ]] && { ((slo++)); continue; }
        ((total++))
        run_one "$isl" "$osl" "$conc"
        case $? in
          0) ((ok++)) ;;
          1) ((fail++)) ;;
          2) ((oom++)); combo_oom=1 ;;
          3) ((slo++)); combo_slo=1 ;;
        esac
        sleep 2
      done

      local elapsed=$(($(date +%s) - start_time))
      local eta="N/A"
      if (( ok > 0 )); then
        local rate=$(echo "scale=3; $ok/$elapsed" | bc 2>/dev/null || echo "0")
        if [[ "$rate" != "0" ]]; then
          local eta_s=$(echo "scale=0; ($total-$ok)/$rate" | bc 2>/dev/null || echo "0")
          eta="$(date -u -d @${eta_s} +%Hh%Mm 2>/dev/null || echo '?')"
        fi
      fi
      log "PROGRESS ok=${ok}/${total} fail=${fail} oom=${oom} slo=${slo} elapsed=$(echo "scale=1;$elapsed/3600"|bc)h ETA=${eta}"
    done
  done
  done

  local total_h=$(echo "scale=1; $(($(date +%s)-start_time))/3600" | bc)
  log "=== COMPLETE: total=${total} ok=${ok} fail=${fail} oom=${oom} slo=${slo} skip=${skip} time=${total_h}h ==="
  log "Results: ${RESULT_DIR}"
}

main