#!/bin/bash
# setup-common.sh — per-node base install for Federated Grid Fault Detection.
#
# Creates a venv at ~/fl-venv, installs PyTorch via tools/setup-torch.sh (the one
# place to change GPU/CUDA/driver/torch — see that file), then the rest of the deps
# from requirements.txt. GPU is the default profile; setup-torch.sh auto-detects an
# NVIDIA GPU so a mis-scheduled node (GPU requested but none attached) still runs on
# CPU.
#
# Robust against the classic FABRIC first-boot failure: cloud-init / unattended
# upgrades holding the apt lock (see setup-lib.sh). Verifies the venv actually got
# created (exits non-zero if not, so the caller sees the failure).
#
# Usage:  bash setup-common.sh [true|false|auto] [cuda_tag]
#
# Install is configurable via positional args OR FL_* environment variables (env is
# the extensible path the notebook uses); all are consumed by setup-torch.sh:
#   FL_USE_GPU auto|true|false ($1) | FL_CUDA_TAG cu121|cu124|cu126|cpu ($2)
#   FL_NVIDIA_DRIVER 535|550|... | FL_TORCH_SPEC torch|torch==2.5.1|...
set -uo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/setup-lib.sh"

USE_GPU_HINT="${1:-${FL_USE_GPU:-auto}}"
CUDA_TAG="${2:-${FL_CUDA_TAG:-cu121}}"
VENV="$HOME/fl-venv"
PY="$VENV/bin/python"
PIP="$VENV/bin/pip"

echo "### PROGRESS: [setup-common] apt base packages (best effort)"
apt_retry update -qq || echo "### PROGRESS: [setup-common] apt update failed (venv fallback will handle it)"
apt_retry install -y -qq python3-venv python3-pip iproute2 pciutils >/dev/null \
    || echo "### PROGRESS: [setup-common] apt install failed (venv fallback will handle it)"

echo "### PROGRESS: [setup-common] creating venv at $VENV"
rm -rf "$VENV"
if ! python3 -m venv "$VENV" >/dev/null 2>&1 || [ ! -x "$PY" ]; then
    # apt is locked / python3-venv missing — build the venv without apt.
    echo "### PROGRESS: [setup-common] venv module unavailable; bootstrapping virtualenv (no apt)"
    rm -rf "$VENV"
    curl -fsSL https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py || {
        echo "ERROR: cannot fetch get-pip.py"; exit 1; }
    python3 /tmp/get-pip.py --user -q || { echo "ERROR: get-pip failed"; exit 1; }
    python3 -m pip install --user -q virtualenv || { echo "ERROR: pip install virtualenv failed"; exit 1; }
    python3 -m virtualenv "$VENV" || { echo "ERROR: virtualenv failed"; exit 1; }
fi
[ -x "$PY" ] || { echo "ERROR: venv python not found at $PY"; exit 1; }
"$PIP" install --upgrade -q pip wheel || { echo "ERROR: pip upgrade failed"; exit 1; }

# PyTorch + NVIDIA driver — delegated to the dedicated, editable script.
# CUDA_TAG/driver/torch config flows through as FL_* env vars (exported below so the
# child sees them even when they arrived as positional args).
echo "### PROGRESS: [setup-common] installing torch via setup-torch.sh"
FL_VENV="$VENV" FL_USE_GPU="$USE_GPU_HINT" FL_CUDA_TAG="$CUDA_TAG" \
    bash "$SCRIPT_DIR/setup-torch.sh" || { echo "ERROR: setup-torch.sh failed"; exit 1; }

# Everything except torch is declared in requirements.txt. Add deps there, not here.
# §3 uploads it to $HOME; fall back to an inline list if it's absent.
REQ=""
for cand in "$HOME/requirements.txt" "$SCRIPT_DIR/../requirements.txt" "$SCRIPT_DIR/requirements.txt"; do
    [ -f "$cand" ] && REQ="$cand" && break
done
if [ -n "$REQ" ]; then
    echo "### PROGRESS: [setup-common] installing requirements from $REQ"
    "$PIP" install -q -r "$REQ" || { echo "ERROR: requirements.txt install failed"; exit 1; }
else
    echo "### PROGRESS: [setup-common] requirements.txt not found; installing core deps inline"
    "$PIP" install -q "numpy" "flwr>=1.7,<2.0" "pandas" "pyarrow" "kagglehub" \
        || { echo "ERROR: core deps install failed"; exit 1; }
fi

echo "### PROGRESS: [setup-common] verifying torch + flwr"
"$PY" -c "import torch, flwr, numpy; \
print('torch', torch.__version__, 'cuda_available', torch.cuda.is_available(), \
'flwr', flwr.__version__)" || { echo "ERROR: torch/flwr import failed"; exit 1; }

echo "### PROGRESS: [setup-common] done"
