#!/usr/bin/env bash
# SUM Parts - PHASE A: everything that does not need the GPU
#
# Splitting the work here is deliberate. Compiling the CUDA extensions needs
# nvcc, not a GPU, and downloading 18 GB of data needs neither. Those are the
# slow parts (~50 min + ~30 min), so they can run while the card is busy with
# someone else's job.
#
#   preflight -> bootstrap -> conda env -> cuda extensions -> patches
#             -> verify imports -> download data -> splits -> link
#
# Ends by writing SETUP_DONE. run_train.sh refuses to start without it.
#
# The one thing this cannot do for you is accept the HuggingFace dataset gate:
# it needs a browser and a logged-in account. The gate IS per-account, so a
# token from a machine that already accepted it is enough. Preflight fails
# immediately if the token is missing or rejected.
#
# Usage:
#   bash scripts/run_setup.sh
#   bash scripts/run_setup.sh --detach
#   RUN_DRYRUN=1 bash scripts/run_setup.sh     # preflight only
set -uo pipefail

SCRIPTS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUT="$HOME/sum-parts/runs/setup"
LOG="$OUT/setup.log"
STATUS="$OUT/STATUS"
DONE_MARKER="$OUT/SETUP_DONE"

# The extensions are compiled for this architecture. We may not be able to ask
# the GPU which one it is (it can be busy or absent), so it is stated instead.
#   RTX 3060 / 3070 / 3080 / 3090 = 8.6   (Ampere consumer)
#   RTX 4090                      = 8.9
#   A100                          = 8.0
export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.6}"

if [ "${1:-}" = "--detach" ]; then
    mkdir -p "$OUT"
    echo "detaching; log: $LOG"
    setsid nohup bash "${BASH_SOURCE[0]}" > "$OUT/nohup.out" 2>&1 < /dev/null &
    sleep 2
    pgrep -af "run_setup.sh" | grep -v detach || true
    exit 0
fi

mkdir -p "$OUT"
exec > >(tee -a "$LOG") 2>&1

PHASE="starting"
STARTED=$(date '+%F %T')

say()   { echo "[$(date '+%F %T')] $*"; }
head_() { echo; echo "════ $* ════"; }

write_status() {
    {
        echo "state    : $1"
        echo "phase    : $PHASE"
        echo "arch     : $TORCH_CUDA_ARCH_LIST"
        echo "started  : $STARTED"
        echo "updated  : $(date '+%F %T')"
        [ -n "${EXTRA:-}" ] && echo "note     : $EXTRA"
        echo "log      : $LOG"
    } > "$STATUS"
}

die() {
    say "GIVING UP in phase '$PHASE': $*"
    EXTRA="$*" write_status "FAILED"
    exit 1
}

STEP_RETRIES="${STEP_RETRIES:-4}"
STEP_BACKOFF="${STEP_BACKOFF:-60}"

# Every phase is idempotent, so retrying one - or rerunning the whole script -
# is always safe. Existing installs, applied patches and cached archives are
# detected and skipped.
step() {
    PHASE="$1"; shift
    head_ "$PHASE"
    write_status "running"
    local attempt=1 wait=$STEP_BACKOFF
    while :; do
        if "$@"; then
            [ "$attempt" -gt 1 ] && say "phase '$PHASE' succeeded on attempt $attempt"
            return 0
        fi
        [ "$attempt" -ge "$STEP_RETRIES" ] && die "$* (failed $attempt times)"
        say "phase '$PHASE' failed (attempt $attempt/$STEP_RETRIES); retrying in ${wait}s"
        EXTRA="retrying $PHASE ($attempt/$STEP_RETRIES)" write_status "retrying"
        sleep "$wait"
        attempt=$((attempt + 1)); wait=$((wait * 2))
    done
}

step_once() {
    PHASE="$1"; shift
    head_ "$PHASE"
    write_status "running"
    "$@" || die "$*"
}

# ---------------------------------------------------------------- preflight

preflight() {
    local fail=0

    say "checking tools"
    for t in git curl python3 tar; do
        command -v "$t" > /dev/null || { say "  MISSING: $t"; fail=1; }
    done

    say "checking disk (need 35 GB free in \$HOME)"
    local free_gb
    free_gb=$(df -BG --output=avail "$HOME" | tail -1 | tr -dc '0-9')
    say "  free: ${free_gb} GB"
    [ "${free_gb:-0}" -lt 35 ] && { say "  NOT ENOUGH"; fail=1; }

    # A GPU is NOT required for this phase. Report what is there, but never
    # fail on it -- the whole point of the split is to work while the card is
    # occupied.
    say "checking GPU (informational only for this phase)"
    if command -v nvidia-smi > /dev/null; then
        nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv,noheader \
            | sed 's/^/  /' || say "  nvidia-smi failed; continuing anyway"
    else
        say "  nvidia-smi absent -- fine for setup, required before training"
    fi
    say "  building extensions for arch $TORCH_CUDA_ARCH_LIST"

    say "checking HuggingFace credentials"
    local tok=""
    [ -n "${HF_TOKEN:-}" ] && tok="$HF_TOKEN"
    [ -z "$tok" ] && [ -f "$HOME/.cache/huggingface/token" ] \
        && tok=$(tr -d '\r\n' < "$HOME/.cache/huggingface/token")
    if [ -z "$tok" ]; then
        say "  MISSING: no HF token"
        say "  fix: copy it from a machine that already accepted the gate:"
        say "       mkdir -p ~/.cache/huggingface"
        say "       echo hf_xxxxx > ~/.cache/huggingface/token"
        say "  the gate needs a browser once, per account:"
        say "       https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts"
        fail=1
    else
        local code
        code=$(curl -s -o /dev/null -w '%{http_code}' -I \
            -H "Authorization: Bearer $tok" \
            "https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts/resolve/main/demo.zip")
        say "  gate probe: HTTP $code"
        case "$code" in
            200|302) say "  gate OK" ;;
            401|403) say "  REJECTED -- token invalid, or this account has not accepted the gate"
                     fail=1 ;;
            *)       say "  unexpected response; continuing but the download may fail" ;;
        esac
    fi

    [ "$fail" -eq 0 ] || return 1
    say "preflight OK -- no GPU needed from here to the end of this phase"
}

do_verify() {
    source "$HOME/miniconda3/etc/profile.d/conda.sh"
    conda activate sumparts
    # SKIP_CUDA_CHECK keeps this to imports only. Querying the device is
    # harmless even on a busy card, but it fails outright if the driver is not
    # present yet - and that must not block a GPU-free setup.
    SKIP_CUDA_CHECK=1 WANDB_MODE=disabled python "$SCRIPTS/verify_env.py"
}

# ---------------------------------------------------------------- run

say "run_setup starting (phase A - no GPU required)"
say "log: $LOG"
write_status "running"

step_once "preflight"   preflight

if [ -n "${RUN_DRYRUN:-}" ]; then
    say "DRYRUN set -- preflight passed, stopping before the real work"
    write_status "dryrun-ok"
    exit 0
fi

step "bootstrap"        bash "$SCRIPTS/bootstrap.sh"
step "conda env"        bash "$SCRIPTS/setup_env.sh"
step "cuda extensions"  bash "$SCRIPTS/setup_pointnext.sh"
step "patch numpy"      bash "$SCRIPTS/patch_numpy_aliases.sh"
step "patch test split" bash "$SCRIPTS/patch_unlabeled_test.sh"
step "patch val mode"   bash "$SCRIPTS/patch_val_mode.sh"
step "verify imports"   do_verify
step "download data"    bash "$SCRIPTS/download_data.sh" all
step "prepare splits"   bash "$SCRIPTS/prepare_full_split.sh"
step "link data"        bash "$SCRIPTS/link_data.sh"

PHASE="done"
touch "$DONE_MARKER"
write_status "DONE"

head_ "PHASE A COMPLETE"
say "built for arch : $TORCH_CUDA_ARCH_LIST"
say "marker         : $DONE_MARKER"
echo
say "next, once the GPU is free:"
say "  bash scripts/run_train.sh"
say "or, to keep it alive unattended:"
say "  setsid nohup bash scripts/keepalive.sh run_train > ~/keepalive.out 2>&1 &"