#!/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"
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

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"
}

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}"

  set +e
  uv pip install --index-url "$TORCH_INDEX" "${TORCH_SPEC[@]}"
  TORCH_RC=$?
  set -e

  # 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"
    )
    uv pip install --index-url "$TORCH_INDEX" "${TORCH_SPEC[@]}"
  elif [[ "$TORCH_RC" -ne 0 ]]; then
    echo ""
    echo "ERROR: PyTorch install failed."
    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
    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 [[ -d "$REPO_DIR" && -f "$REPO_DIR/venv/bin/activate" && -d "$REPO_DIR/ui" ]]; then
  start_existing_install
  exit 0
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