#!/bin/bash
set -euo pipefail

REPO_BASE="https://docs.grid.tf/threefold_public/public/raw/branch/master"

echo "=== Grid VM Preparation ==="

if [ "$EUID" -ne 0 ]; then
  echo "Please run as root"
  exit 1
fi

# --- System update and base tools ---
echo ""
echo "=== Updating system and installing base tools ==="
apt-get update -y
DEBIAN_FRONTEND=noninteractive apt-get install -y \
  sudo nmon tmux restic tcpdump nano iputils-ping net-tools curl wget

# --- TF Users ---
echo ""
echo "=== Setting up TF Users ==="
wget -q "${REPO_BASE}/add-tf-users.sh" -O /tmp/add-tf-users.sh
bash /tmp/add-tf-users.sh
rm -f /tmp/add-tf-users.sh

# --- SSH hardening ---
echo ""
echo "=== Configuring SSH ==="

ubuntu_version=$(lsb_release -rs 2>/dev/null || (. /etc/os-release && echo "$VERSION_ID"))
major_version=$(echo "$ubuntu_version" | cut -d. -f1)

echo "Detected Ubuntu $ubuntu_version (major: $major_version)"

# Backup original if not already backed up
if [ ! -f "/etc/ssh/sshd_config.original" ]; then
  cp /etc/ssh/sshd_config /etc/ssh/sshd_config.original
  echo "[OK] Backed up original sshd_config"
fi

wget -q "${REPO_BASE}/sshd_config" -O /etc/ssh/sshd_config
echo "[OK] Installed hardened sshd_config (port 34022, no root, no password)"

if [ "$major_version" -ge 24 ]; then
  # Ubuntu 24.04+ uses ssh.socket — override listen port
  echo "Configuring ssh.socket for port 34022..."
  mkdir -p /etc/systemd/system/ssh.socket.d
  cat > /etc/systemd/system/ssh.socket.d/port.conf <<EOF
[Socket]
ListenStream=
ListenStream=34022
EOF
  systemctl daemon-reload
  systemctl restart ssh.socket
  systemctl enable ssh.socket
  echo "[OK] ssh.socket configured for port 34022"
else
  # Ubuntu 22.04 and older — restart ssh service
  systemctl restart ssh
  echo "[OK] ssh service restarted"
fi

# --- Optional flags ---
while getopts ":dc" opt; do
  case ${opt} in
    d )
      echo ""
      echo "=== Installing Docker ==="
      if command -v docker &>/dev/null; then
        echo "[OK] Docker already installed"
      else
        DEBIAN_FRONTEND=noninteractive apt-get install -y \
          ca-certificates curl gnupg lsb-release
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
        echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
        apt-get update
        DEBIAN_FRONTEND=noninteractive apt-get install -y docker-ce docker-ce-cli containerd.io
        systemctl enable --now docker
        echo "[OK] Docker installed"
      fi
      ;;
    c )
      echo ""
      echo "=== Installing Caddy ==="
      if command -v caddy &>/dev/null; then
        echo "[OK] Caddy already installed"
      else
        apt-get install -y debian-keyring debian-archive-keyring apt-transport-https
        curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' | gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg
        curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' | tee /etc/apt/sources.list.d/caddy-stable.list
        apt-get update
        DEBIAN_FRONTEND=noninteractive apt-get install -y caddy
        echo "[OK] Caddy installed"
      fi
      ;;
    \? )
      echo "Invalid option: $OPTARG" 1>&2
      ;;
  esac
done
shift $((OPTIND -1))

echo ""
echo "=== VM Preparation Complete ==="
echo ""
echo "SSH is now on port 34022. Connect with:"
echo "  ssh -p 34022 <username>@<server-ip>"
echo ""
echo "Root login is disabled. Use sudo from a TF user account."