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

export DEBIAN_FRONTEND=noninteractive

if command -v sudo >/dev/null 2>&1; then
  SUDO=sudo
else
  SUDO=
fi

mkdir -p /workspace/hf-cache /workspace/datasets /workspace/checkpoints /workspace/reports /workspace/tmp

if command -v apt-get >/dev/null 2>&1; then
  $SUDO apt-get update
  $SUDO apt-get install -y --no-install-recommends \
    git \
    git-lfs \
    curl \
    ca-certificates \
    build-essential \
    cmake \
    ninja-build \
    tmux \
    htop \
    rsync \
    jq \
    p7zip-full \
    libaio-dev
fi

git lfs install || true

cat >/tmp/infosec-training-env.sh <<'EOF'
export HF_HOME=/workspace/hf-cache
export HF_HUB_CACHE=/workspace/hf-cache/hub
export TRANSFORMERS_CACHE=/workspace/hf-cache/transformers
export DATASETS_CACHE=/workspace/hf-cache/datasets
export WANDB_PROJECT=infosec-qwen-cybergym
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
EOF

if [[ -w /etc/profile.d ]]; then
  cp /tmp/infosec-training-env.sh /etc/profile.d/infosec-training.sh
fi

set +u
source /tmp/infosec-training-env.sh
set -u

python -m pip install --upgrade pip setuptools wheel packaging ninja

if [[ -f training/requirements-cu126.txt ]]; then
  python -m pip install -r training/requirements-cu126.txt
else
  echo "training/requirements-cu126.txt not found yet; clone the repo and rerun this script for Python deps."
fi

python - <<'PY'
import os
import sys

print("Python:", sys.version)
print("HF_HOME:", os.environ.get("HF_HOME"))
try:
    import torch
    print("Torch:", torch.__version__)
    print("CUDA available:", torch.cuda.is_available())
    if torch.cuda.is_available():
        print("CUDA devices:", torch.cuda.device_count())
        for i in range(torch.cuda.device_count()):
            print(i, torch.cuda.get_device_name(i))
except Exception as exc:
    print("Torch check failed:", repr(exc))
PY

echo "Bootstrap complete."