#!/usr/bin/env bash
# Build a non-destructive, self-contained archive of the MiniMax-H3 experiments.
set -Eeuo pipefail

ROOT=${SSKJ_H3_ROOT:-/data/wxy/sskj-h3}
SOURCE_MAP="$ROOT/SOURCE_MAP.tsv"

log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }

command -v rsync >/dev/null || die "rsync is required"
command -v sha256sum >/dev/null || die "sha256sum is required"

for section in sglang-base vllm-omni-base sglang-profile vbench-base; do
  mkdir -p "$ROOT/$section/scripts" "$ROOT/$section/results"
done
mkdir -p "$ROOT/vbench-base/environment" "$ROOT/tools"

copy_file() {
  local source=$1 destination=$2
  [[ -f "$source" ]] || die "missing source file: $source"
  mkdir -p "$(dirname "$destination")"
  rsync -a "$source" "$destination"
}

copy_tree() {
  local source=$1 destination=$2
  [[ -d "$source" ]] || die "missing source directory: $source"
  mkdir -p "$destination"
  rsync -a --exclude='__pycache__/' --exclude='*.pyc' "$source/" "$destination/"
}

log "copying SGLang base scripts and results"
copy_file /data/wxy/run_sglang_h3_mixed_matrix_6000d.sh "$ROOT/sglang-base/scripts/run_sglang_h3_mixed_matrix_6000d.sh"
copy_file /data/wxy/minimax_h3_mixed_bench.py "$ROOT/sglang-base/scripts/minimax_h3_mixed_bench.py"
copy_tree /data/wxy/results/minimax_h3_mixed_matrix/mixed64-20steps-5s-20260822-100844 "$ROOT/sglang-base/results/mixed64-20steps-5s-20260822-100844"
copy_tree /data/wxy/results/minimax_h3_mixed_matrix/balanced-tp4-tp2-20steps-5s-20260822-175030 "$ROOT/sglang-base/results/balanced-tp4-tp2-20steps-5s-20260822-175030"

log "copying vLLM-Omni base scripts and results"
copy_file /data/wxy/run_vllm_omni_h3_matrix_6000d.sh "$ROOT/vllm-omni-base/scripts/run_vllm_omni_h3_matrix_6000d.sh"
copy_file /data/wxy/minimax_h3_vllm_bench.py "$ROOT/vllm-omni-base/scripts/minimax_h3_vllm_bench.py"
copy_tree /data/wxy/results/minimax_h3_vllm_matrix/vllm-balanced64-20steps-5s-r2-20260822-225317 "$ROOT/vllm-omni-base/results/vllm-balanced64-20steps-5s-r2-20260822-225317"

log "copying SGLang profile scripts and results"
copy_tree /data/wxy/h3_profile "$ROOT/sglang-profile/scripts/h3_profile"
copy_tree /data/wxy/profile_results/h3-quick-input-matrix-20260824-run1 "$ROOT/sglang-profile/results/h3-quick-input-matrix-20260824-run1"
copy_tree /data/wxy/profile_results/h3-targeted-profile-20260824-run1 "$ROOT/sglang-profile/results/h3-targeted-profile-20260824-run1"
copy_tree /data/wxy/profile_results/h3-sdpa-kernel-matrix-20260825-run1 "$ROOT/sglang-profile/results/h3-sdpa-kernel-matrix-20260825-run1"

log "copying VBench scripts, scoring environment evidence, and results"
for name in \
  check_vbench_score_env.py h3_vbench_generate.py prefetch_vbench_weights.sh \
  rebuild_vbench_score_env.sh run_h3_vbench_base_tp2x4.sh setup_vbench_6000d.sh; do
  copy_file "/data/wxy/$name" "$ROOT/vbench-base/scripts/$name"
done
copy_file /data/wxy/h3_vbench_base_README.md "$ROOT/vbench-base/ORIGINAL_RUNBOOK.md"
copy_tree /data/wxy/vbench_score_setup_logs "$ROOT/vbench-base/environment/vbench_score_setup_logs"
copy_file /data/wxy/vbench_score_env_rebuild.log "$ROOT/vbench-base/environment/vbench_score_env_rebuild.log"
copy_file /data/wxy/vbench_score_env_red.log "$ROOT/vbench-base/environment/vbench_score_env_red.log"
copy_tree /data/wxy/results/h3_vbench_base/h3-vbench-base-dense-tp2x4-20260826-run1 "$ROOT/vbench-base/results/h3-vbench-base-dense-tp2x4-20260826-run1"

copy_file /data/wxy/organize_sskj_h3.sh "$ROOT/tools/organize_sskj_h3.sh"
copy_file /data/wxy/validate_sskj_h3_archive.py "$ROOT/tools/validate_sskj_h3_archive.py"

log "building source map"
printf 'section\tkind\tsource\tdestination\tfiles\tbytes\n' > "$SOURCE_MAP"

append_map() {
  local section=$1 kind=$2 source=$3 destination=$4 files bytes
  if [[ -f "$source" ]]; then
    files=1
    bytes=$(stat -c '%s' "$source")
  else
    read -r files bytes < <(
      find "$source" -type f ! -path '*/__pycache__/*' ! -name '*.pyc' -printf '%s\n' |
        awk '{n+=1; b+=$1} END {print n+0, b+0}'
    )
  fi
  printf '%s\t%s\t%s\t%s\t%s\t%s\n' \
    "$section" "$kind" "$source" "$destination" "$files" "$bytes" >> "$SOURCE_MAP"
}

append_map sglang-base script /data/wxy/run_sglang_h3_mixed_matrix_6000d.sh sglang-base/scripts/run_sglang_h3_mixed_matrix_6000d.sh
append_map sglang-base script /data/wxy/minimax_h3_mixed_bench.py sglang-base/scripts/minimax_h3_mixed_bench.py
append_map sglang-base result /data/wxy/results/minimax_h3_mixed_matrix/mixed64-20steps-5s-20260822-100844 sglang-base/results/mixed64-20steps-5s-20260822-100844
append_map sglang-base result /data/wxy/results/minimax_h3_mixed_matrix/balanced-tp4-tp2-20steps-5s-20260822-175030 sglang-base/results/balanced-tp4-tp2-20steps-5s-20260822-175030
append_map vllm-omni-base script /data/wxy/run_vllm_omni_h3_matrix_6000d.sh vllm-omni-base/scripts/run_vllm_omni_h3_matrix_6000d.sh
append_map vllm-omni-base script /data/wxy/minimax_h3_vllm_bench.py vllm-omni-base/scripts/minimax_h3_vllm_bench.py
append_map vllm-omni-base result /data/wxy/results/minimax_h3_vllm_matrix/vllm-balanced64-20steps-5s-r2-20260822-225317 vllm-omni-base/results/vllm-balanced64-20steps-5s-r2-20260822-225317
append_map sglang-profile script /data/wxy/h3_profile sglang-profile/scripts/h3_profile
append_map sglang-profile result /data/wxy/profile_results/h3-quick-input-matrix-20260824-run1 sglang-profile/results/h3-quick-input-matrix-20260824-run1
append_map sglang-profile result /data/wxy/profile_results/h3-targeted-profile-20260824-run1 sglang-profile/results/h3-targeted-profile-20260824-run1
append_map sglang-profile result /data/wxy/profile_results/h3-sdpa-kernel-matrix-20260825-run1 sglang-profile/results/h3-sdpa-kernel-matrix-20260825-run1
append_map vbench-base result /data/wxy/results/h3_vbench_base/h3-vbench-base-dense-tp2x4-20260826-run1 vbench-base/results/h3-vbench-base-dense-tp2x4-20260826-run1
append_map vbench-base environment /data/wxy/vbench_score_setup_logs vbench-base/environment/vbench_score_setup_logs

log "building script checksum manifests"
for section in sglang-base vllm-omni-base sglang-profile vbench-base; do
  (
    cd "$ROOT/$section"
    find scripts -type f ! -name SHA256SUMS -print0 | sort -z | xargs -0 sha256sum > scripts/SHA256SUMS
  )
done

log "archive data copy complete: $ROOT"