#!/bin/bash
set -euo pipefail

# Check if AUTH_PASS environment variable was provided (safely handles unbound variables)
if [ -z "${AUTH_PASS:-}" ]; then
  echo "ERROR: Password environment variable is missing!"
  echo "Usage: AUTH_PASS='your_password' curl -fsSL <url> | bash"
  exit 1
fi

# Install directly in the SSH user's home directory
INSTALL_BASE="${HOME}"
AI_TOOLKIT_DIR="${INSTALL_BASE}/ai-toolkit"
UI_PORT=8675

SETUP_LOG="/tmp/ai_toolkit_setup.log"
UI_LOG="/tmp/ai_toolkit_ui.log"

: > "${SETUP_LOG}"

# Color formatting
if [ -t 1 ]; then
  RESET=$'\033[0m'
  BOLD=$'\033[1m'
  RED=$'\033[31m'
  GREEN=$'\033[32m'
  YELLOW=$'\033[33m'
  BLUE=$'\033[34m'
  CYAN=$'\033[36m'
else
  RESET=""
  BOLD=""
  RED=""
  GREEN=""
  YELLOW=""
  BLUE=""
  CYAN=""
fi

log() { printf "%b[%s]%b %s\n" "${BLUE}" "INFO" "${RESET}" "$1"; }
err() { printf "%b[%s]%b %s\n" "${RED}" "ERROR" "${RESET}" "$1" >&2; exit 1; }
ok()  { printf "%b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$1"; }
section() { printf "\n%b%s%b\n" "${BOLD}${CYAN}" "$1" "${RESET}"; }
kv()  { printf "  %b%-22s%b %s\n" "${YELLOW}" "$1" "${RESET}" "$2"; }

run_logged() {
  local label="$1"
  local logfile="$2"
  shift 2

  printf "  %b[%s]%b %s\n" "${BLUE}" ".." "${RESET}" "$label"
  if "$@" >>"${logfile}" 2>&1; then
    printf "  %b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$label"
  else
    printf "  %b[%s]%b %s\n" "${RED}" "FAIL" "${RESET}" "$label" >&2
    printf "%s\n" "--- last 30 lines from ${logfile} ---" >&2
    tail -n 30 "${logfile}" >&2 || true
    exit 1
  fi
}

# ------------------ 1. WORKSPACE PREPARATION ------------------
section "1. Preparing Workspace Directory"
mkdir -p "${INSTALL_BASE}"
ok "Installation home directory: ${INSTALL_BASE}"

kv "Install Base (Home)" "${INSTALL_BASE}"
kv "ai-toolkit Path"     "${AI_TOOLKIT_DIR}"
kv "Setup Log"           "${SETUP_LOG}"

# ------------------ 2. SYSTEM DEPENDENCIES & NVM NODE.JS 24 ------------------
section "2. Installing System Dependencies & Node.js 24 via NVM"

run_logged "apt update" "${SETUP_LOG}" sudo apt-get update -y || apt-get update -y
run_logged "apt packages" "${SETUP_LOG}" sudo apt-get install -y \
  curl git screen python3 python3-venv python3-pip \
  zip unzip nano ffmpeg build-essential || apt-get install -y \
  curl git screen python3 python3-venv python3-pip \
  zip unzip nano ffmpeg build-essential

# Download and install nvm & Node.js 24
log "Installing nvm and Node.js 24..."
export NVM_DIR="$HOME/.nvm"
if [ ! -s "$NVM_DIR/nvm.sh" ]; then
  run_logged "download nvm" "${SETUP_LOG}" bash -lc 'curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.6/install.sh | bash'
fi

# Load nvm into script context
\. "$HOME/.nvm/nvm.sh"

run_logged "install Node.js 24" "${SETUP_LOG}" nvm install 24
nvm use 24

log "Node.js version: $(node -v)"
log "npm version: $(npm -v)"

# Install uv for fast Python package installation
if ! command -v uv >/dev/null 2>&1; then
  run_logged "install uv" "${SETUP_LOG}" bash -lc 'curl -Ls https://astral.sh/uv/install.sh | sh'
fi

export PATH="${HOME}/.local/bin:${HOME}/.cargo/bin:/root/.cargo/bin:$PATH"
if [ -f "${HOME}/.local/bin/env" ]; then
  source "${HOME}/.local/bin/env" || true
fi

ok "System dependencies & Node.js 24 installed"

# ------------------ 3. AI-TOOLKIT REPO & PYTHON VENV ------------------
section "3. Cloning ai-toolkit & Setting up Environment"

cd "${INSTALL_BASE}"

if [ ! -d "${AI_TOOLKIT_DIR}" ]; then
  run_logged "git clone ai-toolkit" "${SETUP_LOG}" git clone --recursive https://github.com/ostris/ai-toolkit.git "${AI_TOOLKIT_DIR}"
else
  log "ai-toolkit directory exists. Fetching updates..."
  cd "${AI_TOOLKIT_DIR}"
  git pull || true
fi

cd "${AI_TOOLKIT_DIR}"
run_logged "git submodules update" "${SETUP_LOG}" git submodule update --init --recursive

# Create Virtual Environment
if [ ! -d "venv" ]; then
  run_logged "create venv" "${SETUP_LOG}" python3 -m venv venv
fi

source venv/bin/activate
run_logged "upgrade pip" "${SETUP_LOG}" pip install --upgrade pip

# Determine installer (uv if available, else pip)
INSTALLER="pip install"
if command -v uv >/dev/null 2>&1; then
  INSTALLER="uv pip install"
elif [ -f "${HOME}/.local/bin/uv" ]; then
  INSTALLER="${HOME}/.local/bin/uv pip install"
fi

# Install PyTorch & Dependencies using specified CUDA 12.8 wheel index
log "Installing PyTorch 2.11.0 (CUDA 12.8)..."
run_logged "install torch cu128" "${SETUP_LOG}" ${INSTALLER} torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu128

if [ -f "requirements.txt" ]; then
  run_logged "install requirements.txt" "${SETUP_LOG}" ${INSTALLER} -r requirements.txt
fi

if [ -f "setup.py" ] || [ -f "pyproject.toml" ]; then
  run_logged "install ai-toolkit package" "${SETUP_LOG}" ${INSTALLER} -e . || true
fi

run_logged "install HF hub & utilities" "${SETUP_LOG}" ${INSTALLER} \
  huggingface_hub[cli] hf_transfer tensorboard wandb bitsandbytes diffusers transformers peft accelerate

# Create standard subdirectories in ai-toolkit
mkdir -p dataset config models output

ok "Python environment setup complete"

# ------------------ 4. AI-TOOLKIT WEB UI SETUP ------------------
section "4. Setting up AI-Toolkit Web UI"

if [ -d "${AI_TOOLKIT_DIR}/ui" ]; then
  cd "${AI_TOOLKIT_DIR}/ui"
  run_logged "npm install (ui dependencies)" "${SETUP_LOG}" npm install
  ok "Web UI dependencies installed"
else
  err "ui directory not found in ai-toolkit."
fi

# ------------------ 5. LAUNCH WEB UI SCREEN SESSION ------------------
section "5. Launching Web UI"

screen -dmS aitoolkit_ui bash -lc "
  export NVM_DIR=\"\$HOME/.nvm\"
  [ -s \"\$NVM_DIR/nvm.sh\" ] && \. \"\$NVM_DIR/nvm.sh\"
  nvm use 24
  cd ${AI_TOOLKIT_DIR}/ui
  source ../venv/bin/activate
  export HF_HUB_ENABLE_HF_TRANSFER=1
  echo '=== STARTING AI-TOOLKIT WEB UI (Port ${UI_PORT}) ==='
  AI_TOOLKIT_AUTH=$AUTH_PASS npm run build_and_start > ${UI_LOG} 2>&1
  exec bash
"
log "Screen session 'aitoolkit_ui' launched on port ${UI_PORT}"

# ------------------ SUMMARY ------------------
section "Setup Complete!"
kv "ai-toolkit Path" "${AI_TOOLKIT_DIR}"
kv "Node.js Version" "$(node -v)"
kv "PyTorch Version" "2.11.0 (cu128)"
kv "Web UI Port"     "${UI_PORT}"
kv "Screen Session"  "screen -r aitoolkit_ui"

section "🌐 Direct Connection Info"
echo "  • Access URL: http://<YOUR_VM_IP>:${UI_PORT}"
echo "  • Local/SSH:  http://localhost:${UI_PORT}"
echo "  • Attach UI:  screen -r aitoolkit_ui"
echo ""
ok "Done!"