#!/usr/bin/env bash
# Legion host bootstrap — run ONCE on a fresh Ubuntu 24.04 install, before Terraform.
#
# Usage (from Mac after restoring ~/Secrets):
#   ssh-copy-id legion
#   scp ~/Secrets/cloudflared-gitea.json legion:/tmp/cloudflared-gitea.json
#   ssh legion "bash -s" < bootstrap.sh
#
# After this script: run `direnv exec . terraform apply` from ~/Development/infrastructure/servers/legion/

set -euo pipefail

LEGION_IP="${1:-192.168.8.148}"
NVIDIA_DRIVER="nvidia-driver-580"

log() { echo ""; echo "==> $*"; }

# ── 0. Sanity check ──────────────────────────────────────────────────────────
if [ "$(id -u)" != "0" ]; then
  echo "Run as root or via sudo. Re-running with sudo..."
  exec sudo -E bash "$0" "$@"
fi

log "Starting Legion bootstrap (IP: ${LEGION_IP})"

# ── 1. Base packages ──────────────────────────────────────────────────────────
log "Installing base prerequisites"
apt-get update -qq
apt-get install -y --no-install-recommends \
  curl wget gnupg ca-certificates apt-transport-https \
  nfs-common open-iscsi \
  jq unzip git

# Enable open-iscsi (required for some k3s storage backends)
systemctl enable --now iscsid || true

# ── 2. Disable systemd-resolved (conflicts with AdGuard on port 53) ───────────
log "Disabling systemd-resolved"
systemctl disable --now systemd-resolved || true
rm -f /etc/resolv.conf
echo "nameserver 1.1.1.1" > /etc/resolv.conf
echo "nameserver 8.8.8.8" >> /etc/resolv.conf
chattr +i /etc/resolv.conf  # prevent NetworkManager from overwriting it

# ── 3. NVIDIA drivers ─────────────────────────────────────────────────────────
log "Installing NVIDIA drivers (${NVIDIA_DRIVER})"
add-apt-repository -y ppa:graphics-drivers/ppa
apt-get update -qq
apt-get install -y "${NVIDIA_DRIVER}"

# NVIDIA container toolkit (required for k3s GPU support)
log "Installing NVIDIA container toolkit"
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
  | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L 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://#g' \
  > /etc/apt/sources.list.d/nvidia-container-toolkit.list
apt-get update -qq
apt-get install -y nvidia-container-toolkit

# Configure NVIDIA runtime for containerd (k3s uses containerd)
nvidia-ctk runtime configure --runtime=containerd

# ── 4. k3s ────────────────────────────────────────────────────────────────────
log "Configuring k3s"
mkdir -p /etc/rancher/k3s
cat > /etc/rancher/k3s/config.yaml << EOF
write-kubeconfig-mode: "644"
default-runtime: nvidia
tls-san:
  - ${LEGION_IP}
  - 127.0.0.1
EOF

log "Installing k3s"
curl -sfL https://get.k3s.io | sh -

log "Waiting for k3s to start..."
sleep 20
k3s kubectl get nodes

# Copy kubeconfig to user home (for 'will' user)
mkdir -p /home/will/.kube
cp /etc/rancher/k3s/k3s.yaml /home/will/.kube/config
sed -i "s/127.0.0.1/${LEGION_IP}/g" /home/will/.kube/config
chown will:will /home/will/.kube/config

# ── 5. NVIDIA device plugin for k3s ──────────────────────────────────────────
log "Installing NVIDIA k3s device plugin"
k3s kubectl create -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.5/nvidia-device-plugin.yml \
  || k3s kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.5/nvidia-device-plugin.yml

# ── 6. Cloudflared (host-level — Gitea SSH tunnel only) ───────────────────────
# The platform tunnel (neuralplatform.ai) runs as a k8s pod managed by Terraform.
# This host-level cloudflared serves only the Gitea SSH NodePort (git.neuralplatform.ai → :30022).
log "Installing cloudflared"
curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg \
  | gpg --dearmor -o /usr/share/keyrings/cloudflare-main.gpg
echo 'deb [signed-by=/usr/share/keyrings/cloudflare-main.gpg] https://pkg.cloudflare.com/cloudflared any main' \
  > /etc/apt/sources.list.d/cloudflared.list
apt-get update -qq
apt-get install -y cloudflared

# Restore tunnel credentials from Mac-side secret (must exist at /tmp/cloudflared-gitea.json)
TUNNEL_ID="b1bf80bf-5448-4c33-8667-d6fce8a82c7a"
if [ -f "/tmp/cloudflared-gitea.json" ]; then
  mkdir -p /etc/cloudflared
  cp /tmp/cloudflared-gitea.json "/etc/cloudflared/${TUNNEL_ID}.json"
  chmod 600 "/etc/cloudflared/${TUNNEL_ID}.json"
  cat > /etc/cloudflared/config.yml << EOF
tunnel: ${TUNNEL_ID}
credentials-file: /etc/cloudflared/${TUNNEL_ID}.json

ingress:
  - hostname: git.neuralplatform.ai
    service: http://localhost:30322
  - service: http_status:404
EOF
  cloudflared service install
  systemctl enable --now cloudflared
  log "Cloudflared installed and running"
else
  log "WARNING: /tmp/cloudflared-gitea.json not found — skipping cloudflared setup"
  log "         To fix: scp ~/Secrets/certs/cloudflared-gitea.json legion:/tmp/cloudflared-gitea.json"
  log "         Then re-run: cloudflared service install && systemctl enable --now cloudflared"
fi

# ── 7. CLI tools (Terraform, Helm) ────────────────────────────────────────────
log "Installing Terraform"
wget -qO- https://apt.releases.hashicorp.com/gpg \
  | gpg --dearmor -o /usr/share/keyrings/hashicorp-archive-keyring.gpg
echo "deb [signed-by=/usr/share/keyrings/hashicorp-archive-keyring.gpg] https://apt.releases.hashicorp.com $(lsb_release -cs) main" \
  > /etc/apt/sources.list.d/hashicorp.list
apt-get update -qq && apt-get install -y terraform

log "Installing Helm"
curl -fsSL https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash

# ── 8. Done ───────────────────────────────────────────────────────────────────
log "Bootstrap complete!"
echo ""
echo "Next steps:"
echo "  1. Copy kubeconfig to Mac:  scp legion:/home/will/.kube/config ~/.kube/legion-config"
echo "  2. Apply Terraform:         cd ~/Development/infrastructure/servers/legion && direnv exec . terraform apply"
echo "  3. Restore Gitea data:      see RUNBOOK.md § Data Restore"
echo ""
echo "Verify GPU access:"
echo "  ssh legion kubectl describe node | grep nvidia.com/gpu"