#!/usr/bin/env bash
# Bootstrap script voor vast.ai NGC PyTorch base image.
# Draait bij instance-start, installeert Cineleum-deps + code + start training.
#
# Path: this file wordt gedownload op de instance via een 1-liner in launcher;
# dan kickt het de hele training-run af.
#
# Env vars expected (via --env flag in vastai create):
#   HF_TOKEN (required, write-scope)
#   HF_DATASET_REPO      default: HaaDeej/cineleum-training-data
#   HF_CHECKPOINT_REPO   default: HaaDeej/cineleum-controlnet-sdxl
#   DATASET_KEY          default: archive_sdxl_v1
#   CHECKPOINT_PREFIX    default: run_a_baseline
#   (+ training hyperparams: LR, MAX_TRAIN_STEPS, etc — passed to handler_train.sh)

set -euo pipefail

LOG_PREFIX="[bootstrap]"
echo "$LOG_PREFIX starting on $(hostname) at $(date -Iseconds)"

: "${HF_TOKEN:?HF_TOKEN env var required}"
: "${HF_DATASET_REPO:=HaaDeej/cineleum-training-data}"
: "${HF_CHECKPOINT_REPO:=HaaDeej/cineleum-controlnet-sdxl}"
: "${CINELEUM_CODE_REPO:=HaaDeej/cineleum-controlnet-sdxl}"
: "${CINELEUM_CODE_BRANCH:=main}"

export HF_TOKEN
export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"

# System deps (NGC image has PyTorch/CUDA but not all our Python deps)
# Uninstall NGC-bundled flash-attn + xformers first — they have ABI mismatch
# with the specific torch patch version, causing diffusers import failures.
# SDXL-ControlNet training @ 512px batch=4 op H100 80GB heeft geen xformers nodig.
echo "$LOG_PREFIX uninstalling ABI-broken flash-attn/xformers..."
pip uninstall -y flash-attn xformers 2>&1 | tail -3 || true

echo "$LOG_PREFIX installing python deps..."
pip install --no-cache-dir --quiet \
    "diffusers==0.31.0" \
    "transformers==4.46.3" \
    "accelerate==1.0.1" \
    "peft==0.13.2" \
    "safetensors==0.4.3" \
    "bitsandbytes==0.49.2" \
    "huggingface-hub==0.25.0" \
    "pyiqa==0.1.15.post2" \
    "scenedetect==0.6.7.1" \
    "ImageHash==4.3.2" \
    "opencv-python-headless==4.11.0.86" \
    "lpips==0.1.4" \
    "numpy==1.26.4" \
    "pillow==12.2.0" \
    "scipy==1.17.1" \
    "tqdm==4.67.1" \
    "wandb" 2>&1 | tail -20 || {
        echo "$LOG_PREFIX pip install FAILED" ; exit 1
    }
echo "$LOG_PREFIX deps installed"

# Fetch Cineleum code (training scripts + handler) from HF model repo.
# We stash source-code tarballs daarin bij elke release.
mkdir -p /workspace/scripts /workspace/code
cd /workspace

echo "$LOG_PREFIX downloading code tarball..."
huggingface-cli download "$CINELEUM_CODE_REPO" \
    "code/cineleum_code.tar.gz" \
    --repo-type=model --local-dir=/tmp/cineleum_code \
    --token="$HF_TOKEN" 2>&1 | tail -5
tar xzf /tmp/cineleum_code/code/cineleum_code.tar.gz -C /workspace/
rm -rf /tmp/cineleum_code

# Sanity
for f in handler_train.sh scripts/train_controlnet_sdxl.py; do
    if [ ! -f "/workspace/$f" ]; then
        echo "$LOG_PREFIX ERROR: missing /workspace/$f after code-extract"
        exit 2
    fi
done
echo "$LOG_PREFIX code ready in /workspace"

chmod +x /workspace/handler_train.sh

# Hand off to training handler
echo "$LOG_PREFIX handing off to handler_train.sh"
exec /workspace/handler_train.sh