#!/usr/bin/env bash

set -euo pipefail

FAMILY="${1:?usage: run_real_pilot.sh FAMILY RHO TP MNS SERVER_PORT}"
RHO="${2:?usage: run_real_pilot.sh FAMILY RHO TP MNS SERVER_PORT}"
TP="${3:?usage: run_real_pilot.sh FAMILY RHO TP MNS SERVER_PORT}"
MNS="${4:?usage: run_real_pilot.sh FAMILY RHO TP MNS SERVER_PORT}"
SERVER_PORT="${5:?usage: run_real_pilot.sh FAMILY RHO TP MNS SERVER_PORT}"

case "${TP}" in
  1|2|4) ;;
  *) echo "ERROR: TP must be 1, 2, or 4" >&2; exit 2 ;;
esac
case "${MNS}" in
  8|16|32|64) ;;
  *) echo "ERROR: MNS must be 8, 16, 32, or 64" >&2; exit 2 ;;
esac

EXPERIMENT_ROOT="${EXPERIMENT_ROOT:-/home/admin/cpfs/wjh/aituner/workload-regime-20260720}"
TRACES_MANIFEST="${TRACES_MANIFEST:-${EXPERIMENT_ROOT}/traces/manifest.json}"
VENV_ROOT="${VENV_ROOT:-/home/admin/cpfs/wjh/venvs/vllm-0.20.0-cu129-workload-regime-v2}"
RUNNER_ROOT="${RUNNER_ROOT:-/home/admin/cpfs/wjh/aituner/aituner-graph-piecewise-bdc357d/runs/frontier-fidelity-envelope-v1}"
RUNNER="${RUNNER_ROOT}/run_qwen30_exact_trace_real_anchor.sh"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
OUTPUT_GENERATION="${OUTPUT_GENERATION:-v1}"

for path in "${TRACES_MANIFEST}" "${VENV_ROOT}/bin/vllm" "${RUNNER}" \
  "${MODEL_ROOT}/config.json"; do
  [[ -e "${path}" ]] || { echo "ERROR: missing ${path}" >&2; exit 1; }
done

mapfile -t CASE_FIELDS < <(
  jq -r --arg family "${FAMILY}" --argjson rho "${RHO}" '
    .cases
    | map(select(.family == $family and ((.rho - $rho) | fabs) < 1e-12))
    | .[]
    | [.private_jsonl, (.prefix_caching | tostring), .private_jsonl_sha256]
    | @tsv
  ' "${TRACES_MANIFEST}"
)
[[ "${#CASE_FIELDS[@]}" -eq 1 ]] || {
  echo "ERROR: expected one manifest case for family=${FAMILY} rho=${RHO}, got ${#CASE_FIELDS[@]}" >&2
  exit 1
}
IFS=$'\t' read -r REQUESTS_FILE PREFIX_CACHING EXPECTED_REQUESTS_SHA \
  <<< "${CASE_FIELDS[0]}"
[[ -f "${REQUESTS_FILE}" ]] || { echo "ERROR: missing ${REQUESTS_FILE}" >&2; exit 1; }
ACTUAL_REQUESTS_SHA="$(sha256sum "${REQUESTS_FILE}" | cut -d' ' -f1)"
[[ "${ACTUAL_REQUESTS_SHA}" == "${EXPECTED_REQUESTS_SHA}" ]] || {
  echo "ERROR: private trace digest mismatch" >&2
  exit 1
}

RHO_LABEL="${RHO//./p}"
TRACE_LABEL="${FAMILY}-rho${RHO_LABEL}-tp${TP}-mns${MNS}"
OUTPUT_ROOT="${EXPERIMENT_ROOT}/real-pilot-${OUTPUT_GENERATION}/${TRACE_LABEL}"
[[ ! -e "${OUTPUT_ROOT}/results/result.json" ]] || {
  echo "ERROR: refusing to overwrite completed ${OUTPUT_ROOT}" >&2
  exit 1
}

FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-/tmp/wjh/flashinfer-workload-regime-vllm020-cu129}"
mkdir -p "${FLASHINFER_WORKSPACE_BASE}"
echo "WORKLOAD_REGIME_REAL_PILOT family=${FAMILY} rho=${RHO} tp=${TP} mns=${MNS} prefix=${PREFIX_CACHING} requests_sha256=${ACTUAL_REQUESTS_SHA} venv=${VENV_ROOT} output=${OUTPUT_ROOT}"

cd "${RUNNER_ROOT}"
OUTPUT_ROOT="${OUTPUT_ROOT}" REQUESTS_FILE="${REQUESTS_FILE}" \
TP="${TP}" MNS="${MNS}" TRACE_LABEL="${TRACE_LABEL}" \
SERVER_PORT="${SERVER_PORT}" PREFIX_CACHING="${PREFIX_CACHING}" \
VENV_ROOT="${VENV_ROOT}" MODEL_ROOT="${MODEL_ROOT}" \
FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE}" \
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-300}" \
bash "${RUNNER}"