#!/usr/bin/env bash
# Runs on the GPU pod. Reads HF_TOKEN and GH_TOKEN from stdin (one per line), then:
#  1. Installs uv + vLLM (pinned to a recent stable that supports Gemma-3-27B + dynamic LoRA)
#  2. Downloads the Gemma-3-27B-it base model to /workspace/models
#  3. Downloads the davidafrica/functional-wellbeing LoRA subfolder to /workspace/adapter
#  4. Installs the local working tree of this repo (scp'd to /workspace/code) so we can run analyses on-pod
#  5. Starts vLLM with --enable-lora and the FT LoRA registered, in the background under setsid
#
# All output goes to /workspace/bootstrap.log so we can tail from local.
set -euo pipefail

# 1. Read secrets from env (caller piped them in over ssh stdin and exported them).
[ -z "${HF_TOKEN:-}" ] && { echo "no HF_TOKEN in env"; exit 1; }
[ -z "${GH_TOKEN:-}" ] && { echo "no GH_TOKEN in env"; exit 1; }
export HF_HUB_TOKEN="$HF_TOKEN"

mkdir -p /workspace/models /workspace/adapter /workspace/code
cd /workspace

LOG=/workspace/bootstrap.log
exec > >(tee -a "$LOG") 2>&1
echo "[bootstrap] starting at $(date -u +%FT%TZ)"

# 2. Install uv
if ! command -v uv >/dev/null 2>&1; then
  echo "[bootstrap] installing uv"
  curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.local/bin:$PATH"
fi
echo "[bootstrap] uv: $(uv --version)"

# 3. Make a venv for the analysis side and install our local project + aligne
cd /workspace/code
if [ ! -d .venv ]; then
  uv venv --python 3.12
fi
. .venv/bin/activate

# Install aligne via gh token + credential helper from the skill recipe.
git config --global credential.helper "!f() { echo username=x-access-token; echo password=$GH_TOKEN; }; f"
# uv pip install our project (which depends on aligne via git+https). uv will pick up
# the git credential helper for github.com.
uv pip install -e .
git config --global --unset credential.helper
unset GH_TOKEN

echo "[bootstrap] aligne imports?"
python -c "import aligne; from aligne.metrics.preferences import run_panel, PanelConfig; from aligne.client import ChatClient, Endpoint; print('aligne OK')"

# 4. Install vLLM in its OWN venv (different deps; pin to a recent stable).
echo "[bootstrap] installing vllm in /workspace/vllm-venv"
if [ ! -d /workspace/vllm-venv ]; then
  uv venv --python 3.12 /workspace/vllm-venv
fi
uv pip install --python /workspace/vllm-venv/bin/python "vllm>=0.10.2,<0.12" "huggingface_hub>=0.34"

# 5. Download Gemma-3-27B-it (gated) and the FT adapter in parallel.
echo "[bootstrap] launching model + adapter downloads"
HF_PY=/workspace/vllm-venv/bin/python
$HF_PY -c "
import os
from huggingface_hub import snapshot_download
print('downloading base model...')
p = snapshot_download(
    repo_id='google/gemma-3-27b-it',
    local_dir='/workspace/models/gemma-3-27b-it',
    token=os.environ['HF_TOKEN'],
    allow_patterns=['*.json','*.safetensors','*.txt','*.model','tokenizer*'],
)
print('base ready at', p)
" &
BASE_PID=$!

$HF_PY -c "
import os
from huggingface_hub import snapshot_download
print('downloading FT adapter subfolder...')
p = snapshot_download(
    repo_id='davidafrica/functional-wellbeing',
    local_dir='/workspace/adapter',
    token=os.environ['HF_TOKEN'],
    allow_patterns=['checkpoints/gemma-3-27b_step325/*'],
)
print('adapter snapshot at', p)
" &
ADAPTER_PID=$!

wait $BASE_PID
wait $ADAPTER_PID
echo "[bootstrap] downloads done"

# Sanity check disk:
df -h /workspace

# 6. Launch vLLM in the background (setsid so it survives ssh drop).
ADAPTER_PATH=/workspace/adapter/checkpoints/gemma-3-27b_step325
ls -lh "$ADAPTER_PATH"

setsid /workspace/vllm-venv/bin/vllm serve /workspace/models/gemma-3-27b-it \
    --served-model-name google/gemma-3-27b-it \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 2 \
    --lora-modules "functional-wellbeing=$ADAPTER_PATH" \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.90 \
    --dtype bfloat16 \
    --port 8000 \
    --host 0.0.0.0 \
    > /workspace/vllm.log 2>&1 < /dev/null &
VLLM_PID=$!
echo "[bootstrap] vLLM PID=$VLLM_PID; tailing /workspace/vllm.log"

echo "[bootstrap] done initial setup at $(date -u +%FT%TZ)"
echo "[bootstrap] hint: tail -f /workspace/vllm.log on the pod, or curl http://localhost:8000/v1/models"