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

commit=d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5
repo_url=https://github.com/thu-ml/SageAttention.git
root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
patch="$root/research/sage2_temporal_pair/patches/0001-sm89-temporal-two-pair-qk-pv.patch"
validator="$root/research/sage2_temporal_pair/validate_short.py"
compose="$root/compose.spark.yml"

work_root=${SAGE2_TEMPORAL_WORK_ROOT:-/home/daniel/aeon-spark-test/h3/sage2-temporal-pair}
run_id=${SAGE2_TEMPORAL_RUN_ID:-$(date -u +%Y%m%dT%H%M%SZ)}
run_root="$work_root/$run_id"
repo="$run_root/repo"
baseline_source="$run_root/baseline-source"
candidate_source="$run_root/candidate-source"
builds="$run_root/builds"
output_host=/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks
output_container=/output/h3-blackwell-runtime/benchmarks

if [[ -e "$run_root" ]]; then
  echo "run directory already exists: $run_root" >&2
  exit 2
fi
mkdir -p "$run_root" "$builds/baseline" "$builds/temporal-pair" "$output_host"

git clone --filter=blob:none --no-checkout "$repo_url" "$repo"
git -C "$repo" worktree add --detach "$baseline_source" "$commit"
git -C "$repo" worktree add --detach "$candidate_source" "$commit"
git -C "$candidate_source" apply --check "$patch"
git -C "$candidate_source" apply "$patch"

build_variant() {
  local name=$1
  local source=$2
  docker compose -f "$compose" run --rm --no-deps \
    -e TORCH_CUDA_ARCH_LIST=12.1 \
    -e MAX_JOBS="${MAX_JOBS:-8}" \
    -v "$source:/work/sageattention:rw" \
    -v "$builds:/variants:rw" \
    --workdir /work/sageattention \
    h3-blackwell-runtime \
    bash -lc "set -o pipefail; python setup.py build_ext --force --build-temp /variants/$name/temp --build-lib /variants/$name/lib 2>&1 | tee /variants/$name/build.log"
}

build_variant baseline "$baseline_source"
build_variant temporal-pair "$candidate_source"

baseline_matches=("$builds"/baseline/lib/sageattention/_qattn_sm89*.so)
candidate_matches=("$builds"/temporal-pair/lib/sageattention/_qattn_sm89*.so)
if [[ ${#baseline_matches[@]} -ne 1 || ! -f "${baseline_matches[0]}" ]]; then
  echo "expected one baseline _qattn_sm89 extension" >&2
  exit 3
fi
if [[ ${#candidate_matches[@]} -ne 1 || ! -f "${candidate_matches[0]}" ]]; then
  echo "expected one candidate _qattn_sm89 extension" >&2
  exit 3
fi
extension=$(basename "${baseline_matches[0]}")
if [[ $(basename "${candidate_matches[0]}") != "$extension" ]]; then
  echo "baseline and candidate extension names differ" >&2
  exit 3
fi

short_json="$output_container/gb10-sage2-temporal-pair-short-$run_id.json"
timing_json="$output_container/gb10-sage2-temporal-pair-timing-$run_id.json"

if [[ ${RUN_SANITIZER:-1} == 1 ]]; then
  docker compose -f "$compose" run --rm --no-deps \
    -v "$builds:/variants:ro" \
    -v "$validator:/research/validate_short.py:ro" \
    h3-blackwell-runtime \
    compute-sanitizer --tool memcheck --error-exitcode=99 \
    python /research/validate_short.py \
      --quick \
      --baseline "/variants/baseline/lib/sageattention/$extension" \
      --candidate "/variants/temporal-pair/lib/sageattention/$extension"
fi

docker compose -f "$compose" run --rm --no-deps \
  -v "$builds:/variants:ro" \
  -v "$validator:/research/validate_short.py:ro" \
  h3-blackwell-runtime \
  python /research/validate_short.py \
    --baseline "/variants/baseline/lib/sageattention/$extension" \
    --candidate "/variants/temporal-pair/lib/sageattention/$extension" \
    --output "$short_json"

# The real harness first verifies the deployed baseline SHA, byte-compares both
# variants, then times them in rotating order. No extension is installed.
docker compose -f "$compose" run --rm --no-deps \
  -v "$root:/opt/h3-blackwell-runtime:ro" \
  -v "$builds:/variants:ro" \
  --workdir /opt/h3-blackwell-runtime \
  h3-blackwell-runtime \
  python tools/profile_sage2_register_variants.py \
    --output "$timing_json" \
    --rounds 5 \
    --warmup 3 \
    --iterations 10 \
    --variant "baseline=/variants/baseline/lib/sageattention/$extension" \
    --variant "temporal-pair=/variants/temporal-pair/lib/sageattention/$extension"

timing_host="$output_host/gb10-sage2-temporal-pair-timing-$run_id.json"
python3 - "$timing_host" <<'PY'
import json
import pathlib
import sys

report = json.loads(pathlib.Path(sys.argv[1]).read_text(encoding="utf-8"))
variants = {item["name"]: item for item in report["variants"]}
candidate = variants["temporal-pair"]
latency_ms = candidate["timing"]["p50_s"] * 1000.0
print(f"temporal-pair p50: {latency_ms:.3f} ms")
if not candidate["byte_exact"]:
    raise SystemExit("candidate failed real SHA parity")
if latency_ms >= 220.0:
    raise SystemExit("candidate failed the <220 ms mainloop gate")
PY

if [[ ${RUN_NCU:-0} == 1 ]]; then
  docker compose -f "$compose" run --rm --no-deps \
    --cap-add SYS_ADMIN --cap-add SYS_PTRACE \
    -v "$root:/opt/h3-blackwell-runtime:ro" \
    -v "${candidate_matches[0]}:/opt/venv/lib/python3.12/site-packages/sageattention/$extension:ro" \
    --workdir /opt/h3-blackwell-runtime \
    h3-blackwell-runtime \
    /opt/nvidia/nsight-compute/2025.3.1/ncu \
      --target-processes all \
      --profile-from-start off \
      --section LaunchStats \
      --section Occupancy \
      --section SchedulerStats \
      --section WarpStateStats \
      --section SpeedOfLight \
      --kernel-name regex:qk_int_sv_f8_attn_kernel \
      --launch-count 1 \
      --force-overwrite \
      -o "$output_container/gb10-sage2-temporal-pair-counters-$run_id" \
      python tools/profile_sage2_scheduler.py \
        --output "$output_container/gb10-sage2-temporal-pair-capture-$run_id.json" \
        --cuda-profiler-capture \
        --skip-tail-study
fi

echo "experiment artifacts: $run_root"
echo "short parity: $output_host/gb10-sage2-temporal-pair-short-$run_id.json"
echo "timing: $timing_host"