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

# Where to download gpu-power.service and set-gpu-power.sh
BASE_URL="https://files.hashborn.space/start"

if [[ $EUID -ne 0 ]]; then
  echo "❌ Run as root: sudo $0"
  exit 1
fi

REAL_USER="${SUDO_USER:-root}"

echo "=== System setup (GPU power, Docker, NVIDIA runtime, languages) ==="
echo "User: ${REAL_USER}"

apt-get update
apt-get install -y ca-certificates curl gnupg lsb-release git build-essential

# ---------------- GPU POWER ----------------
echo
read -rp "👉 Configure GPU power limit service (systemd)? [y/N]: " GPU_CONFIRM
if [[ "$GPU_CONFIRM" =~ ^[Yy]$ ]]; then
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "❌ nvidia-smi not found. Install NVIDIA driver + reboot, then re-run stage 2."
    exit 1
  fi

  echo
  echo "Detected GPUs:"
  nvidia-smi -L || true

  echo
  read -rp "Enter Power Limit (Watts), e.g. 160: " GPU_POWER
  if ! [[ "$GPU_POWER" =~ ^[0-9]+$ ]]; then
    echo "❌ Power limit must be an integer number of watts."
    exit 1
  fi

  echo "$GPU_POWER" > /etc/gpu-power-limit.conf
  chmod 644 /etc/gpu-power-limit.conf

  echo "⬇️ Downloading gpu-power.service and set-gpu-power.sh from ${BASE_URL}"
  curl -fsSL "${BASE_URL}/set-gpu-power.sh" -o /usr/local/bin/set-gpu-power.sh
  curl -fsSL "${BASE_URL}/gpu-power.service" -o /etc/systemd/system/gpu-power.service

  chmod 755 /usr/local/bin/set-gpu-power.sh
  chmod 644 /etc/systemd/system/gpu-power.service

  systemctl daemon-reload
  systemctl enable gpu-power.service
  systemctl start gpu-power.service

  echo
  systemctl status gpu-power.service --no-pager || true
fi

# ---------------- DOCKER ----------------
echo
read -rp "👉 Install Docker Engine + Compose v2? [y/N]: " DOCKER_CONFIRM
if [[ "$DOCKER_CONFIRM" =~ ^[Yy]$ ]]; then
  apt-get remove -y docker docker-engine docker.io containerd runc || true

  mkdir -p /etc/apt/keyrings
  curl -fsSL https://download.docker.com/linux/ubuntu/gpg | \
    gpg --dearmor -o /etc/apt/keyrings/docker.gpg

  echo \
"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
https://download.docker.com/linux/ubuntu \
$(lsb_release -cs) stable" \
> /etc/apt/sources.list.d/docker.list

  apt-get update
  apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

  # Optional standalone docker-compose binary (without hyphen command)
  echo
  read -rp "👉 Install standalone /usr/local/bin/docker-compose (no hyphen)? [y/N]: " DC2_CONFIRM
  if [[ "$DC2_CONFIRM" =~ ^[Yy]$ ]]; then
    curl -fsSL \
      "https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m)" \
      -o /usr/local/bin/docker-compose
    chmod +x /usr/local/bin/docker-compose
  fi

  usermod -aG docker "${REAL_USER}" || true

  echo
  echo "Docker versions:"
  docker --version || true
  docker compose version || true
  if command -v docker-compose >/dev/null 2>&1; then
    docker-compose version || true
  fi

  echo
  echo "⚠️ Log out/in (or run: newgrp docker) for 'docker' group to apply."
fi

# ---------------- NVIDIA CONTAINER TOOLKIT ----------------
echo
read -rp "👉 Install NVIDIA Container Toolkit (for Docker GPU runtime)? [y/N]: " NCTK_CONFIRM
if [[ "$NCTK_CONFIRM" =~ ^[Yy]$ ]]; then
  curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

  curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null

  apt-get update
  apt-get install -y nvidia-container-toolkit

  if command -v nvidia-ctk >/dev/null 2>&1; then
    nvidia-ctk runtime configure --runtime=docker
    nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml
  fi

  systemctl restart docker || true
  echo "✅ NVIDIA container toolkit installed."
fi

# ---------------- LANGUAGES ----------------
echo
read -rp "👉 Install languages (Go, Rust, Python dev tools)? [y/N]: " LANG_CONFIRM
if [[ "$LANG_CONFIRM" =~ ^[Yy]$ ]]; then
  apt-get install -y python3 python3-pip python3-venv golang

  # Rust for the real user (not root)
  if [[ "${REAL_USER}" != "root" ]]; then
    sudo -u "${REAL_USER}" bash -lc '
      if ! command -v rustc >/dev/null 2>&1; then
        curl https://sh.rustup.rs -sSf | sh -s -- -y
      fi
      source "$HOME/.cargo/env" >/dev/null 2>&1 || true
      rustc --version || true
      cargo --version || true
    '
  else
    echo "⚠️ REAL_USER=root, skipping rustup install (run it under your user)."
  fi

  echo "Go: $(go version 2>/dev/null || true)"
  echo "Python: $(python3 --version 2>/dev/null || true)"
fi

echo
echo "✅ Stage 2 complete."
echo "If you added user to docker group: log out/in or run: newgrp docker"
