#!/usr/bin/env bash

set -euo pipefail

CAMPAIGN_ROOT="${CAMPAIGN_ROOT:?CAMPAIGN_ROOT is required}"
PROFILE_ROOT="${PROFILE_ROOT:?PROFILE_ROOT is required}"
RUNNER_DIR="${RUNNER_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
FRONTIER_SOURCE="${FRONTIER_SOURCE:-/home/admin/cpfs/wjh/aituner/frontier-q235-v020-5b953f5}"
REPLAYSERVE_ROOT="${REPLAYSERVE_ROOT:-/home/admin/cpfs/wjh/replayserve}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8}"
PYTHON_DEPS="${PYTHON_DEPS:-${VENV_ROOT}/lib/python3.12/site-packages}"
ALLREDUCE_CSV="${ALLREDUCE_CSV:-${RUNNER_DIR}/profiles/measured-allreduce.csv}"
BASE_PUBLIC=/home/admin/cpfs/wjh/aituner/fidelity-envelope-private/trace-exact-v1/public/u0p01/frontier.csv
BASE_PRIVATE=/home/admin/cpfs/wjh/aituner/fidelity-envelope-private/trace-exact-v1/private/u0p01/real_requests.jsonl
FIXED_PER_GPU_RATE="${FIXED_PER_GPU_RATE:-0.2}"
REQUESTS=129

mkdir -p "${CAMPAIGN_ROOT}"/{provenance,traces,real,sim,analysis}
exec > >(tee -a "${CAMPAIGN_ROOT}/controller.log") 2>&1
echo "Q235_CAMPAIGN_LAUNCH_ECHO host=dash0 model=Qwen3-235B-A22B-FP8 engine=vLLM-0.20.0+cu129 simulator=Frontier-5b953f5-piecewise cases={Fixed-PD:4096x256@0.2rps/gpu,Fixed-PO:4096x1@0.2rps/gpu,Trace-PD:u0p01-original-OSL,Trace-PO:u0p01-OSL1} requests=129 transform=t_prime=t/TP configs={TP4/EP1,TP8/EP8}xMNS{64,128} MBT=8192 trials=3 fresh_server=true metrics=mean,p90(TTFT,TPOT-if-PD,E2E) SLO=not_scored expected_wall=10-30h expected_cost=90-220_H20-GPUh profile=${PROFILE_ROOT} output=${CAMPAIGN_ROOT}"
date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
sha256sum "${BASH_SOURCE[0]}" "${BASE_PUBLIC}" "${BASE_PRIVATE}" "${MODEL_ROOT}/config.json" \
  "${PROFILE_ROOT}/manifest.json" "${ALLREDUCE_CSV}" > "${CAMPAIGN_ROOT}/provenance/input.sha256"
git -C "${FRONTIER_SOURCE}" rev-parse HEAD > "${CAMPAIGN_ROOT}/provenance/frontier.commit"
git -C "${RUNNER_DIR}" rev-parse HEAD > "${CAMPAIGN_ROOT}/provenance/aituner.commit"

NORMALIZER="${RUNNER_DIR}/../simulator-tuning-latency-matrix-v0/materialize_qwen30_tp_normalized_trace.py"
MATERIALIZER="${RUNNER_DIR}/prepare_qwen30_latency_case.py"
for tp in 4 8; do
  "${VENV_ROOT}/bin/python" "${NORMALIZER}" --base-public-csv "${BASE_PUBLIC}" \
    --base-private-jsonl "${BASE_PRIVATE}" --tp "${tp}" \
    --output-root "${CAMPAIGN_ROOT}/traces/trace-pd/tp${tp}"
  "${VENV_ROOT}/bin/python" "${MATERIALIZER}" trace \
    --base-public "${CAMPAIGN_ROOT}/traces/trace-pd/tp${tp}/public/frontier.csv" \
    --base-private "${CAMPAIGN_ROOT}/traces/trace-pd/tp${tp}/private/real_requests.jsonl" \
    --output-tokens 1 --tp "${tp}" --output-root "${CAMPAIGN_ROOT}/traces/trace-po/tp${tp}"
  for spec in fixed-pd:256 fixed-po:1; do
    case_name="${spec%%:*}"; osl="${spec##*:}"
    "${VENV_ROOT}/bin/python" "${MATERIALIZER}" fixed --model "${MODEL_ROOT}" \
      --input-tokens 4096 --output-tokens "${osl}" --requests "${REQUESTS}" \
      --per-gpu-rate "${FIXED_PER_GPU_RATE}" --tp "${tp}" \
      --output-root "${CAMPAIGN_ROOT}/traces/${case_name}/tp${tp}"
  done
done

run_real() {
  local case_name="$1" prefix="$2" port="$3"
  CASE_NAME="${case_name}" PREFIX_CACHING="${prefix}" \
    TRACE_ROOT="${CAMPAIGN_ROOT}/traces/${case_name}" OUTPUT_ROOT="${CAMPAIGN_ROOT}/real/${case_name}" \
    RUNNER_DIR="${RUNNER_DIR}" VENV_ROOT="${VENV_ROOT}" MODEL_ROOT="${MODEL_ROOT}" \
    BASE_PORT="${port}" bash "${RUNNER_DIR}/run_qwen235_v020_real_surface.sh"
}
run_real fixed-pd false 9300
run_real fixed-po false 9400
run_real trace-pd true 9800
run_real trace-po true 9900

"${VENV_ROOT}/bin/python" "${RUNNER_DIR}/extract_qwen235_v020_runtime_contract.py" \
  --case-root "${CAMPAIGN_ROOT}/real/fixed-pd" \
  --output "${CAMPAIGN_ROOT}/provenance/runtime-contract.json"

run_sim() {
  local case_name="$1" prefix_flag="$2"
  "${VENV_ROOT}/bin/python" "${RUNNER_DIR}/run_frontier_qwen235_v020_surface.py" \
    --frontier-source "${FRONTIER_SOURCE}" --replayserve-root "${REPLAYSERVE_ROOT}" \
    --profile-root "${PROFILE_ROOT}" --python-deps "${PYTHON_DEPS}" \
    --output-root "${CAMPAIGN_ROOT}/sim/${case_name}" \
    --runtime-contract "${CAMPAIGN_ROOT}/provenance/runtime-contract.json" \
    --trace-tp "4=${CAMPAIGN_ROOT}/traces/${case_name}/tp4/public/frontier.csv" \
    --trace-tp "8=${CAMPAIGN_ROOT}/traces/${case_name}/tp8/public/frontier.csv" \
    "${prefix_flag}" --allreduce-csv "${ALLREDUCE_CSV}" --resume
}
run_sim fixed-pd --no-prefix-caching
run_sim fixed-po --no-prefix-caching
run_sim trace-pd --prefix-caching
run_sim trace-po --prefix-caching

"${VENV_ROOT}/bin/python" "${RUNNER_DIR}/analyze_qwen235_v020_campaign.py" \
  --campaign-root "${CAMPAIGN_ROOT}" \
  --json-output "${CAMPAIGN_ROOT}/analysis/comparison.json" \
  --markdown-output "${CAMPAIGN_ROOT}/analysis/comparison.md"
find "${CAMPAIGN_ROOT}" -type f ! -path '*/provenance/artifacts.sha256' -print0 \
  | sort -z | xargs -0 sha256sum > "${CAMPAIGN_ROOT}/provenance/artifacts.sha256"
date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
echo Q235_V020_CAMPAIGN_COMPLETE