#!/usr/bin/env bash
#
# AI-Toolkit 1-Click Install or Restart for RunPod - FAST VERSION
# ---------------------------------------------------------------
# Main speed improvements:
# - Skips apt-get update/install when packages are already present
# - Uses persistent caches in /workspace for pip, uv, and npm
# - Uses uv for PyTorch and requirements instead of pip where possible
# - Auto-detects Blackwell / RTX 50-series GPUs instead of asking every time
# - Reuses system Node.js 22+ when available, otherwise installs NVM in /workspace
# - Uses npm ci when package-lock.json exists

set -euo pipefail

WORKDIR="/workspace"
REPO_URL="https://github.com/ostris/ai-toolkit.git"
REPO_DIR="$WORKDIR/ai-toolkit"
INSTALL_MARKER="$REPO_DIR/.runpod_fast_install_complete"
NVM_DIR="$WORKDIR/.nvm"
CACHE_DIR="$WORKDIR/.cache"
PIP_CACHE_DIR="$CACHE_DIR/pip"
UV_CACHE_DIR="$CACHE_DIR/uv"
NPM_CACHE_DIR="$CACHE_DIR/npm"

export PIP_CACHE_DIR
export UV_CACHE_DIR
export npm_config_cache="$NPM_CACHE_DIR"
export DEBIAN_FRONTEND=noninteractive

# RunPod download links can occasionally be very slow or pause for more than
# uv's default 30-second read timeout. Give large PyTorch/NVIDIA wheels enough
# time to continue instead of aborting the entire installation.
export UV_HTTP_TIMEOUT="${UV_HTTP_TIMEOUT:-300}"
export UV_HTTP_RETRIES="${UV_HTTP_RETRIES:-5}"
export PIP_DEFAULT_TIMEOUT="${PIP_DEFAULT_TIMEOUT:-600}"
export PIP_RETRIES="${PIP_RETRIES:-20}"

mkdir -p "$WORKDIR" "$CACHE_DIR" "$PIP_CACHE_DIR" "$UV_CACHE_DIR" "$NPM_CACHE_DIR"

if [[ $(id -u) -ne 0 ]]; then
  SUDO=sudo
else
  SUDO=''
fi

log() {
  echo ""
  echo "============================================================"
  echo "$1"
  echo "============================================================"
}

small_log() {
  echo ""
  echo "--- $1"
}

command_exists() {
  command -v "$1" >/dev/null 2>&1
}

package_installed() {
  dpkg -s "$1" >/dev/null 2>&1
}

ensure_apt_packages() {
  local packages=(git python3 python3-venv python3-pip curl ca-certificates)
  local missing=()

  for pkg in "${packages[@]}"; do
    if ! package_installed "$pkg"; then
      missing+=("$pkg")
    fi
  done

  if [[ ${#missing[@]} -eq 0 ]]; then
    small_log "System packages already installed. Skipping apt-get."
    return 0
  fi

  small_log "Installing missing system packages: ${missing[*]}"
  $SUDO apt-get update -y
  $SUDO apt-get install -y --no-install-recommends "${missing[@]}"
}

ensure_build_tools() {
  if package_installed build-essential; then
    small_log "build-essential already installed."
    return 0
  fi

  small_log "Installing build-essential fallback package."
  $SUDO apt-get update -y
  $SUDO apt-get install -y --no-install-recommends build-essential
}

node_major_version() {
  if ! command_exists node; then
    echo "0"
    return 0
  fi

  node -v 2>/dev/null | sed 's/^v//' | cut -d. -f1 || echo "0"
}

ensure_node_22() {
  export NVM_DIR="$NVM_DIR"

  local current_major
  current_major="$(node_major_version)"

  if [[ "$current_major" =~ ^[0-9]+$ ]] && [[ "$current_major" -ge 22 ]] && command_exists npm; then
    small_log "Using existing Node.js $(node --version)."
    return 0
  fi

  if [[ -f "$NVM_DIR/nvm.sh" ]]; then
    small_log "Using existing NVM installation."
    # shellcheck disable=SC1091
    source "$NVM_DIR/nvm.sh"
    nvm install 22
    nvm use 22 >/dev/null
    return 0
  fi

  small_log "Installing NVM and Node.js 22."
  mkdir -p "$NVM_DIR"
  curl -fsSL https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash

  # shellcheck disable=SC1091
  source "$NVM_DIR/nvm.sh"
  nvm install 22
  nvm use 22 >/dev/null
}

detect_cuda_stream() {
  local gpu_name=""

  if command_exists nvidia-smi; then
    gpu_name="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -n 1 || true)"
  fi

  echo "Detected GPU: ${gpu_name:-unknown}"

  # Blackwell / RTX 50-series needs newer CUDA wheels.
  # Examples: RTX 5090, RTX 5080, RTX 5070, RTX PRO 6000 Blackwell, B200, GB200.
  if echo "$gpu_name" | grep -Eiq 'RTX[[:space:]-]*(50|PRO.*Blackwell)|Blackwell|B200|GB200|GB300'; then
    CUDA_STREAM="cu128"
  else
    CUDA_STREAM="cu126"
  fi

  echo "Selected PyTorch CUDA wheel: $CUDA_STREAM"
}

install_torch_with_fallback() {
  local torch_index="$1"
  shift
  local torch_packages=("$@")
  local install_rc=1
  local max_attempts=2

  # First use uv for speed. The extended timeout/retry environment variables
  # above protect large NVIDIA wheels from temporary network stalls.
  for attempt in $(seq 1 "$max_attempts"); do
    echo "PyTorch uv attempt $attempt of $max_attempts..."

    if uv pip install --index-url "$torch_index" "${torch_packages[@]}"; then
      return 0
    else
      install_rc=$?
    fi

    if [[ "$attempt" -lt "$max_attempts" ]]; then
      echo "PyTorch download failed. Retrying in $((attempt * 5)) seconds..."
      sleep $((attempt * 5))
    fi
  done

  echo ""
  echo "uv could not finish the PyTorch download. Falling back to pip."
  echo "PyPI will be used for NVIDIA runtime packages and the PyTorch index for Torch wheels."

  # Keeping PyPI as the primary index avoids relying exclusively on the
  # occasionally slow pypi.nvidia.com links exposed by the PyTorch index.
  if python -m pip install \
    --timeout "$PIP_DEFAULT_TIMEOUT" \
    --retries "$PIP_RETRIES" \
    --prefer-binary \
    --index-url "https://pypi.org/simple" \
    --extra-index-url "$torch_index" \
    "${torch_packages[@]}"; then
    return 0
  else
    install_rc=$?
  fi

  return "$install_rc"
}

clone_or_update_repo() {
  cd "$WORKDIR"

  if [[ -d "$REPO_DIR/.git" ]]; then
    small_log "Existing ai-toolkit source found. Updating it instead of cloning again."
    git -C "$REPO_DIR" fetch --depth=1 origin
    git -C "$REPO_DIR" reset --hard origin/main
    return 0
  fi

  if [[ -d "$REPO_DIR" ]]; then
    small_log "Incomplete ai-toolkit folder found. Removing it before clean clone."
    rm -rf "$REPO_DIR"
  fi

  small_log "Cloning ai-toolkit repository."

  git config --global http.postBuffer 524288000 || true
  git config --global http.version HTTP/1.1 || true

  local clone_success=0
  local max_retries=3

  for attempt in $(seq 1 "$max_retries"); do
    echo "Clone attempt $attempt of $max_retries..."

    if git clone --depth=1 --single-branch "$REPO_URL" "$REPO_DIR"; then
      clone_success=1
      break
    fi

    if [[ "$attempt" -lt "$max_retries" ]]; then
      echo "Clone failed. Retrying in 3 seconds..."
      sleep 3
    fi
  done

  git config --global --unset http.version >/dev/null 2>&1 || true
  git config --global --unset http.postBuffer >/dev/null 2>&1 || true

  if [[ "$clone_success" -ne 1 ]]; then
    echo ""
    echo "ERROR: Failed to clone ai-toolkit after $max_retries attempts."
    echo "Please check the pod network and try again."
    exit 1
  fi
}

create_venv_and_install_python_deps() {
  cd "$REPO_DIR"

  log "Creating Python virtual environment"

  if [[ ! -d "$REPO_DIR/venv" ]]; then
    python3 -m venv venv
  fi

  # shellcheck disable=SC1091
  source "$REPO_DIR/venv/bin/activate"

  small_log "Upgrading pip tools and installing uv."
  python -m pip install -U pip setuptools wheel uv

  detect_cuda_stream

  if [[ "$CUDA_STREAM" == "cu128" ]]; then
    TORCH_SPEC=(
      "torch==2.7.0+cu128"
      "torchvision==0.22.0+cu128"
      "torchaudio==2.7.0+cu128"
    )
  else
    TORCH_SPEC=(
      "torch==2.7.0+cu126"
      "torchvision==0.22.0+cu126"
      "torchaudio==2.7.0+cu126"
    )
  fi

  TORCH_INDEX="https://download.pytorch.org/whl/${CUDA_STREAM}"

  log "Installing PyTorch ${CUDA_STREAM}"

  if install_torch_with_fallback "$TORCH_INDEX" "${TORCH_SPEC[@]}"; then
    TORCH_RC=0
  else
    TORCH_RC=$?
  fi

  # Extra fallback for older CUDA streams, in case the index resolves package names
  # without the local +cuXXX suffix on a specific RunPod image.
  if [[ "$TORCH_RC" -ne 0 && "$CUDA_STREAM" == "cu126" ]]; then
    echo ""
    echo "Torch install with +cu126 package names failed. Trying cu126 index without local suffix..."
    TORCH_SPEC=(
      "torch==2.7.0"
      "torchvision==0.22.0"
      "torchaudio==2.7.0"
    )

    if uv pip install --index-url "$TORCH_INDEX" "${TORCH_SPEC[@]}"; then
      TORCH_RC=0
    else
      TORCH_RC=$?
    fi
  fi

  if [[ "$TORCH_RC" -ne 0 ]]; then
    echo ""
    echo "ERROR: PyTorch install failed after uv retries and the pip fallback."
    exit "$TORCH_RC"
  fi

  CONSTRAINTS_FILE="$REPO_DIR/runpod-fast-constraints.txt"
  cat > "$CONSTRAINTS_FILE" <<EOF_CONSTRAINTS
# Keep the correct RunPod Torch wheel locked during dependency install.
${TORCH_SPEC[0]}
${TORCH_SPEC[1]}
${TORCH_SPEC[2]}
EOF_CONSTRAINTS

  log "Installing AI-Toolkit Python requirements"

  set +e
  uv pip install -r requirements.txt -c "$CONSTRAINTS_FILE"
  UV_RC=$?
  set -e

  if [[ "$UV_RC" -ne 0 ]]; then
    echo ""
    echo "uv install failed with exit code $UV_RC. Installing build tools and retrying once..."
    ensure_build_tools

    set +e
    uv pip install -r requirements.txt -c "$CONSTRAINTS_FILE"
    UV_RC_2=$?
    set -e

    if [[ "$UV_RC_2" -ne 0 ]]; then
      echo ""
      echo "uv retry failed. Falling back to pip."
      python -m pip install -r requirements.txt -c "$CONSTRAINTS_FILE"
    fi
  fi

  echo ""
  echo "Python dependencies installed successfully."
}

install_ui_deps_and_launch() {
  ensure_node_22

  cd "$REPO_DIR/ui"

  log "Installing UI dependencies"

  if [[ -f package-lock.json ]]; then
    npm ci --prefer-offline --no-audit --fund=false
  else
    npm install --prefer-offline --no-audit --fund=false
  fi

  log "Building and starting AI-Toolkit UI"

  npm run build_and_start &
  UI_PID=$!

  sleep 3

  if kill -0 "$UI_PID" 2>/dev/null; then
    touch "$INSTALL_MARKER"

    echo ""
    echo "======================================"
    echo "AI-Toolkit installation complete!"
    echo "======================================"
    echo "UI is running. PID: $UI_PID"
    echo ""
    echo "Next time, run this script again and it will start the UI instantly."
    echo ""
  else
    echo ""
    echo "WARNING: UI process may have failed to start."
    echo "Check the logs above for errors."
    echo ""
  fi
}

start_existing_install() {
  log "Existing AI-Toolkit installation detected"

  # shellcheck disable=SC1091
  source "$REPO_DIR/venv/bin/activate"

  ensure_node_22

  cd "$REPO_DIR/ui"

  echo "Starting AI-Toolkit UI..."
  npm run start
}

##############################################################################
# Fast path: existing install -> start the UI
##############################################################################
if [[ -f "$INSTALL_MARKER" && -d "$REPO_DIR" && -f "$REPO_DIR/venv/bin/activate" && -d "$REPO_DIR/ui" ]]; then
  start_existing_install
  exit 0
fi

# A repository/venv can exist after an interrupted installation. Without the
# completion marker, continue through the install path and repair what is missing.
if [[ -d "$REPO_DIR" && ! -f "$INSTALL_MARKER" ]]; then
  small_log "Incomplete previous installation detected. Resuming setup."
fi

##############################################################################
# Full install path
##############################################################################
log "AI-Toolkit first-time installation - fast mode"

ensure_apt_packages
clone_or_update_repo
create_venv_and_install_python_deps
install_ui_deps_and_launch