#!/bin/bash
set -euo pipefail

# Collects run artifacts for a given SLURM job id into a single tar.gz bundle.
# Usage:
#   slurm/collect_gguf_debug_bundle.sh <JOB_ID> [OUTPUT_DIR]

if [[ $# -lt 1 || $# -gt 2 ]]; then
    echo "Usage: $0 <JOB_ID> [OUTPUT_DIR]"
    exit 2
fi

JOB_ID="$1"
OUTPUT_DIR="${2:-}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd -P)"

if [[ -z "$OUTPUT_DIR" ]]; then
    OUTPUT_DIR="${REPO_DIR}/reports/slurm-debug"
fi

mkdir -p "$OUTPUT_DIR"
WORK_DIR="$(mktemp -d "${TMPDIR:-/tmp}/gguf-eval-debug-${JOB_ID}-XXXXXX")"
TIMESTAMP="$(date +%Y%m%dT%H%M%S)"
BUNDLE_PATH="${OUTPUT_DIR}/gguf-eval-debug-${JOB_ID}-${TIMESTAMP}.tar.gz"

cleanup() {
    rm -rf "$WORK_DIR"
}
trap cleanup EXIT

copy_if_exists() {
    local src dst
    src="$1"
    dst="$2"
    if [[ -f "$src" ]]; then
        cp -f "$src" "$dst"
        echo "[OK] copied: $src"
    else
        echo "[WARN] missing: $src"
    fi
}

echo "Collecting debug bundle for job: $JOB_ID"
echo "Repo:   $REPO_DIR"
echo "Output: $BUNDLE_PATH"

mkdir -p "$WORK_DIR/logs"

copy_if_exists "${REPO_DIR}/eval-${JOB_ID}.out" "$WORK_DIR/logs/eval-${JOB_ID}.out"
if [[ -f "${REPO_DIR}/llm-server-${JOB_ID}.log" ]]; then
    copy_if_exists "${REPO_DIR}/llm-server-${JOB_ID}.log" "$WORK_DIR/logs/llm-server-${JOB_ID}.log"
else
    # Backward compatibility for older runs.
    copy_if_exists "${REPO_DIR}/llama-server-${JOB_ID}.log" "$WORK_DIR/logs/llama-server-${JOB_ID}.log"
fi
copy_if_exists "${REPO_DIR}/enroot-diag-${JOB_ID}.log" "$WORK_DIR/logs/enroot-diag-${JOB_ID}.log"
copy_if_exists "${REPO_DIR}/eval-launch-${JOB_ID}.log" "$WORK_DIR/logs/eval-launch-${JOB_ID}.log"
copy_if_exists "${REPO_DIR}/gpu-util-${JOB_ID}.csv" "$WORK_DIR/logs/gpu-util-${JOB_ID}.csv"
copy_if_exists "${REPO_DIR}/run-config-${JOB_ID}.json" "$WORK_DIR/logs/run-config-${JOB_ID}.json"
copy_if_exists "${REPO_DIR}/failure-dump-${JOB_ID}.json" "$WORK_DIR/logs/failure-dump-${JOB_ID}.json"

if command -v scontrol >/dev/null 2>&1; then
    scontrol show job "$JOB_ID" > "$WORK_DIR/slurm-scontrol-job-${JOB_ID}.txt" || true
else
    echo "[WARN] scontrol not found"
fi

if command -v sacct >/dev/null 2>&1; then
    sacct -j "$JOB_ID" --format=JobID,JobName%30,State,ExitCode,Elapsed,NodeList,AllocTRES \
        > "$WORK_DIR/slurm-sacct-job-${JOB_ID}.txt" || true
else
    echo "[WARN] sacct not found"
fi

cat > "$WORK_DIR/README.txt" <<EOT
Bundle: gguf-eval-debug-${JOB_ID}
Created: $(date -Iseconds)
Repo: ${REPO_DIR}
Files:
- logs/eval-${JOB_ID}.out
- logs/llm-server-${JOB_ID}.log (or logs/llama-server-${JOB_ID}.log for older runs)
- logs/enroot-diag-${JOB_ID}.log
- logs/eval-launch-${JOB_ID}.log
- logs/gpu-util-${JOB_ID}.csv
- logs/run-config-${JOB_ID}.json
- logs/failure-dump-${JOB_ID}.json
- slurm-scontrol-job-${JOB_ID}.txt
- slurm-sacct-job-${JOB_ID}.txt
EOT

tar -czf "$BUNDLE_PATH" -C "$WORK_DIR" .
echo "Bundle created: $BUNDLE_PATH"
