#!/usr/bin/env bash
set -Eeuo pipefail

SCENARIO=${1:?scenario required}
VARIANT=${2:?variant required}
GPU_CSV=${3:?gpu csv required}
REPLICA=${4:?replica required}
PHASE_NAME=${5:?phase name required}

ROOT=/data/wxy
MODEL=/data/hf_models/MiniMax-H3
PYTHON=/root/.miniconda3/envs/sglang/bin/python
SGLANG=/root/.miniconda3/envs/sglang/bin/sglang
CLIENT="$ROOT/h3_profile/h3_profile_client.py"
NSYS=/usr/local/cuda-13.2/bin/nsys
MEDIA_BIN_DIR=/root/.miniconda3/envs/deploy/bin
RESULT_ROOT=/data/wxy/profile_results/h3-targeted-profile-20260824-run1
PHASE="$RESULT_ROOT/phases/$PHASE_NAME"
PORT=$((30010 + REPLICA * 10))
MASTER_PORT=$((31000 + REPLICA * 10))
SCHEDULER_PORT=$((32000 + REPLICA * 10))
SERVER_DIR="$PHASE/server_${REPLICA}_port${PORT}"
CLIENT_DIR="$PHASE/client_${REPLICA}_${SCENARIO}"
LOG="$PHASE/parallel_nsys_orchestrator.log"
SERVER_PID=
declare -a MONITOR_PIDS=()

mkdir -p "$SERVER_DIR/outputs" "$CLIENT_DIR/perf" "$PHASE/system"
log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*" | tee -a "$LOG"; }

cleanup() {
  local rc=$? p
  trap - EXIT INT TERM
  for p in "${MONITOR_PIDS[@]:-}"; do kill -TERM "$p" 2>/dev/null || true; done
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
    kill -INT -- "-$SERVER_PID" 2>/dev/null || kill -INT "$SERVER_PID" 2>/dev/null || true
    for _ in {1..60}; do kill -0 "$SERVER_PID" 2>/dev/null || break; sleep 2; done
    kill -TERM -- "-$SERVER_PID" 2>/dev/null || true
    sleep 2
    kill -KILL -- "-$SERVER_PID" 2>/dev/null || true
  fi
  [[ -n "${SERVER_PID:-}" ]] && wait "$SERVER_PID" 2>/dev/null || true
  if ((rc != 0)); then log "$SCENARIO parallel Nsight failed rc=$rc; resume chain will rerun the missing phase"; fi
  exit "$rc"
}
trap cleanup EXIT INT TERM

[[ ! -f "$PHASE/DONE" ]] || { log 'DONE already exists; nothing to do'; exit 0; }
[[ -x "$NSYS" ]] || { log "nsys missing: $NSYS"; exit 1; }
for check_port in "$PORT" "$((PORT+1))" "$MASTER_PORT" "$SCHEDULER_PORT"; do
  if ss -ltn "sport = :$check_port" | grep -q LISTEN; then log "port occupied: $check_port"; exit 1; fi
done
if nvidia-smi -i "$GPU_CSV" --query-compute-apps=pid --format=csv,noheader | grep -q '[0-9]'; then
  log "GPU $GPU_CSV occupied; refusing to interfere"; exit 1
fi

printf '%s\n' \
  "Concurrent exploratory Nsight: scenario=$SCENARIO GPUs=$GPU_CSV." \
  'Overlaps another Nsight capture and RVA Torch profiling. CUDA/NVTX/NCCL structure is usable; absolute latency, OS runtime, CPU gaps, PCIe and collective timing include cross-workload contention.' \
  > "$PHASE/CONCURRENT_DESIGN_NOTE.txt"

nvidia-smi dmon -s pucvmet -d 1 -o DT > "$PHASE/system/nvidia-dmon.log" 2>&1 & MONITOR_PIDS+=("$!")
pidstat -durh 1 > "$PHASE/system/pidstat.log" 2>&1 & MONITOR_PIDS+=("$!")
iostat -xz 1 > "$PHASE/system/iostat.log" 2>&1 & MONITOR_PIDS+=("$!")
vmstat 1 > "$PHASE/system/vmstat.log" 2>&1 & MONITOR_PIDS+=("$!")

log "starting $VARIANT TP=2 scenario=$SCENARIO GPUs=$GPU_CSV port=$PORT under Nsight Systems"
CUDA_VISIBLE_DEVICES="$GPU_CSV" PATH="$MEDIA_BIN_DIR:$PATH" PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false \
  SGLANG_USE_RUNAI_MODEL_STREAMER=false SGLANG_DIFFUSION_STAGE_LOGGING=1 \
  SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=0 NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=COLL \
  setsid "$NSYS" profile --force-overwrite=true --delay 70 --duration 1800 \
    --trace=cuda,nvtx,nccl,osrt --sample=none --cpuctxsw=none --cuda-graph-trace=node \
    -o "$SERVER_DIR/server" \
    "$SGLANG" serve --model-path "$MODEL" --model-variant "$VARIANT" --backend sglang \
      --performance-mode speed --num-gpus 2 --tp-size 2 --ulysses-degree 1 \
      --use-fsdp-inference false --enable-torch-compile false --batching-max-size 1 \
      --batching-delay-ms 0 --enable-layerwise-nvtx-marker --host 0.0.0.0 --port "$PORT" \
      --master-port "$MASTER_PORT" --scheduler-port "$SCHEDULER_PORT" \
      --output-path "$SERVER_DIR/outputs" --log-requests --log-requests-level 2 \
      > "$SERVER_DIR/server.log" 2>&1 &
SERVER_PID=$!

healthy=0
for _ in {1..360}; do
  if curl -fsS --max-time 5 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then healthy=1; break; fi
  kill -0 "$SERVER_PID" 2>/dev/null || { tail -100 "$SERVER_DIR/server.log"; exit 1; }
  sleep 5
done
((healthy == 1)) || { tail -100 "$SERVER_DIR/server.log"; exit 1; }

log 'server healthy; starting client'
"$PYTHON" "$CLIENT" --deployment "TARGET_NSYS_TP2_PARALLEL_${SCENARIO}" --scenario "$SCENARIO" \
  --port "$PORT" --replica-index "$REPLICA" --model "$MODEL" \
  --reference-image "$ROOT/sskj-MiniMax-H3/assets/reference_images/landscape_mountain_lake.jpg" \
  --reference-images-dir "$ROOT/h3_profile/assets/reference_images_5" \
  --reference-video-1s "$MODEL/assets/r2va.mp4" --reference-video-5s "$MODEL/assets/ref2va.mp4" \
  --reference-video-10s "$MODEL/assets/h3_direct_768p.mp4" --steps 20 --duration 5 \
  --short-edge 768 --aspect-ratio 16:9 --seed 1101 --repeats 1 --warmup 0 --warmup-steps 5 \
  --perf-dir "$CLIENT_DIR/perf" --output "$CLIENT_DIR/results.jsonl" > "$CLIENT_DIR/client.log" 2>&1

log 'client completed; stopping profiled server and finalizing .nsys-rep'
kill -INT -- "-$SERVER_PID" 2>/dev/null || kill -INT "$SERVER_PID" 2>/dev/null || true
for _ in {1..120}; do kill -0 "$SERVER_PID" 2>/dev/null || break; sleep 2; done
if kill -0 "$SERVER_PID" 2>/dev/null; then kill -TERM -- "-$SERVER_PID" 2>/dev/null || true; fi
wait "$SERVER_PID" 2>/dev/null || true
SERVER_PID=

report=
for _ in {1..60}; do
  report=$(find "$SERVER_DIR" -maxdepth 1 -type f -name '*.nsys-rep' -size +0c -print -quit)
  [[ -n "$report" ]] && break
  sleep 2
done
[[ -n "$report" ]] || { log 'missing .nsys-rep after client completion'; exit 1; }
touch "$PHASE/DONE"
log "completed; report=$report"