#!/usr/bin/env bash

set -euo pipefail

RHO="${1:?usage: run_full_real.sh RHO CONFIG TRIAL PORT}"
CONFIG="${2:?usage: run_full_real.sh RHO CONFIG TRIAL PORT}"
TRIAL="${3:?usage: run_full_real.sh RHO CONFIG TRIAL PORT}"
PORT="${4:?usage: run_full_real.sh RHO CONFIG TRIAL PORT}"
CONTROL_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_ROOT="${VENV_ROOT:-/home/admin/cpfs/wjh/venvs/vllm-0.20.0-cu129-workload-regime-v2}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
SERVED_MODEL="qwen3-30b-s3-real-full"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-40960}"
ALLOW_SYNTHETIC_PROMPTS="${ALLOW_SYNTHETIC_PROMPTS:-false}"
ALLOW_LONG_CONTEXT_SERVER="${ALLOW_LONG_CONTEXT_SERVER:-false}"
FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-/tmp/wjh/flashinfer-frontier-s3-real-v0-${CONFIG}}"
SERVER_PID=""
TELEMETRY_PID=""

[[ "${RHO}" =~ ^[A-Za-z0-9._-]+$ ]] || {
  echo "ERROR: rho label contains unsupported characters: ${RHO}" >&2; exit 1;
}
if [[ -n "${TRACE_INPUT_ROOT:-}" ]]; then
  INPUT_NAME=""
else
  case "${RHO}" in
    0p00125) INPUT_NAME="canary-r0p00125" ;;
    0p0025) INPUT_NAME="canary-r0p0025" ;;
    0p005) INPUT_NAME="full-r0p005" ;;
    0p01) INPUT_NAME="full-r0p01" ;;
    0p02) INPUT_NAME="full-r0p02" ;;
    *) echo "ERROR: unsupported built-in rho ${RHO}" >&2; exit 1 ;;
  esac
fi
# Generic tp<N>_mns<M> parser: TP = tensor-parallel size, MNS = max-num-seqs.
if [[ "${CONFIG}" =~ ^tp([0-9]+)_mns([0-9]+)$ ]]; then
  TP="${BASH_REMATCH[1]}"; MNS="${BASH_REMATCH[2]}"
else
  echo "ERROR: unsupported config ${CONFIG}" >&2; exit 1
fi
[[ "${TRIAL}" == "1" || "${TRIAL}" == "2" ]] || {
  echo "ERROR: trial must be 1 or 2" >&2; exit 1;
}
if (( MAX_MODEL_LEN > 40960 )) && [[ "${ALLOW_LONG_CONTEXT_SERVER}" != "true" ]]; then
  echo "ERROR: MAX_MODEL_LEN>40960 requires ALLOW_LONG_CONTEXT_SERVER=true" >&2
  exit 1
fi
if (( MAX_MODEL_LEN > 40960 )); then
  export VLLM_CACHE_ROOT="${VLLM_CACHE_ROOT:-/tmp/wjh/vllm-cache-frontier-s3-real-longctx}"
fi

TRACE_INPUT_ROOT="${TRACE_INPUT_ROOT:-${CONTROL_ROOT}/inputs/${INPUT_NAME}}"
REQUESTS_FILE="${TRACE_INPUT_ROOT}/real_requests.jsonl"
INPUT_MANIFEST="${TRACE_INPUT_ROOT}/manifest.json"
OUTPUT_ROOT="${OUTPUT_ROOT:-${CONTROL_ROOT}/outputs/full-real/rho-${RHO}/${CONFIG}/trial-${TRIAL}}"

for path in "${VENV_ROOT}/bin/python" "${VENV_ROOT}/bin/vllm" \
  "${MODEL_ROOT}/config.json" "${REQUESTS_FILE}" "${INPUT_MANIFEST}" \
  "${CONTROL_ROOT}/qwen30_exact_trace_client.py" \
  "${CONTROL_ROOT}/qwen30_prefill_client.py"; do
  [[ -e "${path}" ]] || { echo "ERROR: missing ${path}" >&2; exit 1; }
done
[[ ! -e "${OUTPUT_ROOT}" ]] || { echo "ERROR: refusing to overwrite ${OUTPUT_ROOT}" >&2; exit 1; }
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" \
  "${OUTPUT_ROOT}/results" "${OUTPUT_ROOT}/metrics" \
  "${OUTPUT_ROOT}/telemetry" "${FLASHINFER_WORKSPACE_BASE}"
exec > >(tee -a "${OUTPUT_ROOT}/logs/controller.log") 2>&1

cleanup() {
  if [[ -n "${TELEMETRY_PID}" ]] && kill -0 "${TELEMETRY_PID}" 2>/dev/null; then
    kill "${TELEMETRY_PID}" 2>/dev/null || true
  fi
  TELEMETRY_PID=""
  if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
    kill -TERM -- "-${SERVER_PID}" 2>/dev/null || true
    for _ in $(seq 1 30); do
      kill -0 "${SERVER_PID}" 2>/dev/null || break
      sleep 1
    done
    kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
  fi
  SERVER_PID=""
}
trap cleanup EXIT INT TERM

IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?fleet GPU allocation is required}"
[[ "${#GPU_IDS[@]}" -eq "${TP}" ]] || {
  echo "ERROR: ${CONFIG} requires exactly ${TP} GPUs" >&2; exit 1;
}
nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu \
  --format=csv,noheader,nounits > "${OUTPUT_ROOT}/provenance/gpus.before.csv"
[[ "$(wc -l < "${OUTPUT_ROOT}/provenance/gpus.before.csv")" -eq 8 ]] || {
  echo "ERROR: expected eight GPUs on host" >&2; exit 1;
}
grep -vq 'NVIDIA H20' "${OUTPUT_ROOT}/provenance/gpus.before.csv" && {
  echo "ERROR: expected eight NVIDIA H20 GPUs" >&2; exit 1;
}
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits \
  > "${OUTPUT_ROOT}/provenance/compute-apps.before.csv" || true
if grep -Eq '^[[:space:]]*[0-9]+' "${OUTPUT_ROOT}/provenance/compute-apps.before.csv"; then
  echo "ERROR: host is no longer fully idle" >&2
  exit 1
fi

"${VENV_ROOT}/bin/python" - <<'PY' > "${OUTPUT_ROOT}/provenance/runtime.versions"
import torch, transformers, vllm
print(f"vllm={vllm.__version__}")
print(f"torch={torch.__version__}")
print(f"torch_cuda={torch.version.cuda}")
print(f"transformers={transformers.__version__}")
if vllm.__version__ != "0.20.0" or torch.version.cuda != "12.9":
    raise SystemExit("runtime version mismatch")
PY
env | sort > "${OUTPUT_ROOT}/provenance/environment.txt"
ps -eo user,pid,ppid,etimes,pcpu,pmem,rss,args --sort=pid \
  > "${OUTPUT_ROOT}/provenance/processes.before.txt"
sha256sum "${REQUESTS_FILE}" "${INPUT_MANIFEST}" \
  "${CONTROL_ROOT}/qwen30_exact_trace_client.py" \
  "${CONTROL_ROOT}/qwen30_prefill_client.py" "${BASH_SOURCE[0]}" \
  > "${OUTPUT_ROOT}/provenance/sources.sha256"
INPUT_MANIFEST="${INPUT_MANIFEST}" REQUESTS_FILE="${REQUESTS_FILE}" \
  MAX_MODEL_LEN="${MAX_MODEL_LEN}" \
  ALLOW_SYNTHETIC_PROMPTS="${ALLOW_SYNTHETIC_PROMPTS}" \
  "${VENV_ROOT}/bin/python" - <<'PY'
import json
import os
from pathlib import Path

manifest = json.loads(Path(os.environ["INPUT_MANIFEST"]).read_text())
source_block_size = int(manifest["source_block_size"])
if source_block_size <= 0 or source_block_size % 16 or manifest["target_block_size"] != 16:
    raise SystemExit("source-to-16 block-size contract failed")
failures = {
    key: value for key, value in manifest["block_contract"].items()
    if key.endswith(("mismatches", "collisions", "conflicts"))
}
if any(failures.values()):
    raise SystemExit(f"block/hash contract failed: {failures}")
if (
    manifest["prompt_contract"]["synthetic_fallback_requests"]
    and os.environ["ALLOW_SYNTHETIC_PROMPTS"].lower() != "true"
):
    raise SystemExit("full real requires real prompt text for every request")
with Path(os.environ["REQUESTS_FILE"]).open() as stream:
    max_request_tokens = max(
        int(row["input_length"]) + int(row["output_length"])
        for row in map(json.loads, stream)
    )
if max_request_tokens > int(os.environ["MAX_MODEL_LEN"]):
    raise SystemExit(
        f"trace max request tokens {max_request_tokens} exceeds "
        f"MAX_MODEL_LEN={os.environ['MAX_MODEL_LEN']}"
    )
PY

export TOKENIZERS_PARALLELISM=false VLLM_USE_V1=1 TORCH_CUDA_ARCH_LIST=9.0 PREFIX_CACHING=true
export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 FLASHINFER_WORKSPACE_BASE
HF_OVERRIDE_ARGS=()
if [[ "${ALLOW_LONG_CONTEXT_SERVER}" == "true" ]]; then
  export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
  HF_OVERRIDE_ARGS=(--hf-overrides "{\"max_position_embeddings\":${MAX_MODEL_LEN}}")
fi
ulimit -n 65536
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
  --host 127.0.0.1 --port "${PORT}" --served-model-name "${SERVED_MODEL}" \
  --tensor-parallel-size "${TP}" --gpu-memory-utilization 0.92 \
  "${HF_OVERRIDE_ARGS[@]}" \
  --max-model-len "${MAX_MODEL_LEN}" --max-num-batched-tokens 8192 --max-num-seqs "${MNS}" \
  --enable-prefix-caching --block-size 16 --enable-chunked-prefill --no-enable-log-requests \
  --enable-logging-iteration-details \
  > "${OUTPUT_ROOT}/logs/server.log" 2>&1 &
SERVER_PID=$!
READY=0
for _ in $(seq 1 900); do
  if curl -fsS --max-time 2 "http://127.0.0.1:${PORT}/v1/models" \
    > "${OUTPUT_ROOT}/results/models.json" 2>/dev/null; then READY=1; break; fi
  kill -0 "${SERVER_PID}" 2>/dev/null || { tail -200 "${OUTPUT_ROOT}/logs/server.log"; exit 1; }
  sleep 3
done
[[ "${READY}" -eq 1 ]] || { echo "ERROR: server readiness timeout" >&2; exit 1; }

"${VENV_ROOT}/bin/python" "${CONTROL_ROOT}/qwen30_prefill_client.py" \
  --port "${PORT}" --served-model "${SERVED_MODEL}" --model-path "${MODEL_ROOT}" \
  --rate 1 --requests 4 --input-tokens 512 --output-tokens 1 \
  --output "${OUTPUT_ROOT}/results/warmup.json"
curl -fsS "http://127.0.0.1:${PORT}/metrics" > "${OUTPUT_ROOT}/metrics/before.prom"
(
  while true; do
    ts=$(date -u +%s.%N)
    snap=$(curl -fsS "http://127.0.0.1:${PORT}/metrics" 2>/dev/null \
      | grep -E 'vllm:(num_requests_running|num_requests_waiting|prefix_cache_hits|prefix_cache_queries|gpu_prefix_cache_hit_rate|gpu_cache_usage)' \
      | tr '\n' '|')
    printf '{"wall_time_epoch_s": %s, "metrics": "%s"}\n' "${ts}" "${snap}" \
      >> "${OUTPUT_ROOT}/telemetry/server-state.jsonl"
    sleep 15
  done
) &
TELEMETRY_PID=$!
"${VENV_ROOT}/bin/python" "${CONTROL_ROOT}/qwen30_exact_trace_client.py" \
  --port "${PORT}" --requests-file "${REQUESTS_FILE}" --served-model "${SERVED_MODEL}" \
  --output "${OUTPUT_ROOT}/results/result.json" --tpot-slo-ms 150 \
  --timeout-seconds 5400
kill "${TELEMETRY_PID}" 2>/dev/null || true
TELEMETRY_PID=""
curl -fsS "http://127.0.0.1:${PORT}/metrics" > "${OUTPUT_ROOT}/metrics/after.prom"
INPUT_MANIFEST="${INPUT_MANIFEST}" OUTPUT_ROOT="${OUTPUT_ROOT}" \
  "${VENV_ROOT}/bin/python" - <<'PY'
import json
import os
from pathlib import Path

manifest = json.loads(Path(os.environ["INPUT_MANIFEST"]).read_text())
result = json.loads((Path(os.environ["OUTPUT_ROOT"]) / "results/result.json").read_text())
if result["contract"]["row_vector_sha256"] != manifest["paired_row_vector_sha256"]:
    raise SystemExit("real/sim paired row-vector digest mismatch")
if result["contract"]["requests"] != manifest["requests"]:
    raise SystemExit("real/sim request-count mismatch")
if not all(request["success"] for request in result["requests"]):
    raise SystemExit("one or more requests failed")
PY

cleanup
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total,memory.used,utilization.gpu \
  --format=csv,noheader,nounits > "${OUTPUT_ROOT}/provenance/gpus.after.csv"
find "${OUTPUT_ROOT}" -type f ! -path '*/provenance/artifacts.sha256' -print0 \
  | sort -z | xargs -0 sha256sum > "${OUTPUT_ROOT}/provenance/artifacts.sha256"
echo "FULL_REAL_COMPLETE rho=${RHO} config=${CONFIG} trial=${TRIAL}"