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

# === Config ===
INSTALL_DIR="${INSTALL_DIR:-/workspace/ai-toolkit}"
UI_PORT="${UI_PORT:-8675}"                        # AI-Toolkit UI default port per docs
AUTH_TOKEN="${AI_TOOLKIT_AUTH:-}"                 # Must be provided via environment
NODE_MAJOR="${NODE_MAJOR:-20}"                    # Install Node.js 20 LTS
PYTHON_BIN="${PYTHON_BIN:-python3}"               # Use the system Python to create venv

if [[ -z "${AUTH_TOKEN}" ]]; then
  echo "ERROR: AI_TOOLKIT_AUTH environment variable is not set."
  echo "Set it in your RunPod environment (e.g., AI_TOOLKIT_AUTH=someStrongPassword) and re-run."
  exit 1
fi

echo "==> Updating apt and installing dependencies..."
apt-get update -y
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
  git curl ca-certificates build-essential pkg-config \
  python3-venv python3-dev \
  ffmpeg libgl1 libglib2.0-0

echo "==> Installing Node.js ${NODE_MAJOR} (required: Node >=18)..."
curl -fsSL https://deb.nodesource.com/setup_${NODE_MAJOR}.x | bash -
DEBIAN_FRONTEND=noninteractive apt-get install -y nodejs

echo "==> Cloning AI-Toolkit into ${INSTALL_DIR}..."
mkdir -p "$(dirname "${INSTALL_DIR}")"
if [[ ! -d "${INSTALL_DIR}" ]]; then
  git clone https://github.com/ostris/ai-toolkit.git "${INSTALL_DIR}"
else
  echo "    Repo already exists, pulling latest..."
  git -C "${INSTALL_DIR}" pull --rebase
fi
git -C "${INSTALL_DIR}" submodule update --init --recursive

echo "==> Creating Python venv and upgrading pip..."
cd "${INSTALL_DIR}"
${PYTHON_BIN} -m venv venv
# shellcheck disable=SC1091
source venv/bin/activate
python -m pip install --upgrade pip wheel setuptools

echo "==> Installing PyTorch 2.8.0 + CUDA 12.8 wheels (torch/vision/audio)..."
# Keep versions aligned; use cu128 wheel index.
pip install --no-cache-dir \
  torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 \
  --index-url https://download.pytorch.org/whl/cu128

echo "==> Installing AI-Toolkit Python requirements..."
pip install --no-cache-dir -r requirements.txt

echo "==> Building and starting the AI-Toolkit UI on port ${UI_PORT}..."
# The UI script uses port 8675 by default; we enforce our env and expose host binding.
export AI_TOOLKIT_AUTH="${AUTH_TOKEN}"
export PORT="${UI_PORT}"
export HOST="0.0.0.0"

cd ui
# Use npm to install/build and then start the UI. This blocks (good for container foreground).
npm install
npm run build_and_start