#!/usr/bin/env bash

set -euo pipefail

export DEBIAN_FRONTEND=noninteractive

API_URL="https://control-1.pike-vimba.ts.net:10000"
INSTALL_SCRIPT_URL="https://control-1.pike-vimba.ts.net"
RESUME_FLAG="/var/tmp/aisdom_install_resume"
PERSISTENT_SCRIPT="/usr/local/bin/aisdom-install.sh"

# ---------------------------------------------------------
# 1. Input Handling & Variable Extraction
# ---------------------------------------------------------
if [ -z "${1:-}" ]; then
    echo "Error: Missing configuration blob."
    exit 1
fi

# Decoded format: NODE_ID|API_TOKEN|TS_KEY|REG_USER|REG_PASS|K3S_URL|K3S_TOKEN
DECODED=$(echo "$1" | base64 --decode)
IFS='|' read -r NODE_ID API_TOKEN TAILSCALE_KEY REG_USER REG_PASS K3S_URL K3S_TOKEN <<< "$DECODED"


send_status() {
    local status="$1"
    local msg="$2"
    echo ">>> [$status] $msg"
    # Note: Using the API token to authenticate status updates
    curl -s -X POST "$API_URL/nodes/$NODE_ID/status" \
        -H "Authorization: Bearer $API_TOKEN" \
        -H "Content-Type: application/json" \
        -d "{\"status\": \"$status\"}" || true
}

# ---------------------------------------------------------
# 2. Reboot & Resume Logic
# ---------------------------------------------------------
if [ -f "$RESUME_FLAG" ]; then
    send_status "RESUMING" "Resuming installation after reboot..."
else
    # Download and save the script to a persistent location
    sudo curl -fsSL -o "$PERSISTENT_SCRIPT" "$INSTALL_SCRIPT_URL"
    sudo chmod +x "$PERSISTENT_SCRIPT"

    send_status "CHECKING_OS" "Verifying Ubuntu version..."
    UBUNTU_VER=$(lsb_release -rs)
    if [[ "$UBUNTU_VER" != "22.04" && "$UBUNTU_VER" != "24.04" ]]; then
        send_status "FAILED" "Error: Only Ubuntu 22.04 and 24.04 are supported."
        exit 1
    fi

    send_status "INSTALLING_DRIVERS" "Installing NVIDIA drivers..."
    sudo apt-get update
    sudo apt-get install -y ubuntu-drivers-common nfs-common
    sudo ubuntu-drivers install

    # Create the resume service
    cat <<EOF | sudo tee /etc/systemd/system/aisdom-resume.service
[Unit]
Description=Resume AISDOM Provider Installation
After=network.target tailscaled.service

[Service]
Type=oneshot
ExecStart=/usr/bin/bash $PERSISTENT_SCRIPT "$1"
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF

    sudo systemctl enable aisdom-resume.service
    sudo touch "$RESUME_FLAG"

    send_status "REBOOTING" "Rebooting to activate drivers..."
    sudo reboot
    exit 0
fi

# Cleanup resume logic
sudo systemctl disable aisdom-resume.service || true
sudo rm -f /etc/systemd/system/aisdom-resume.service || true

# ---------------------------------------------------------
# 3. NVIDIA Container Toolkit Setup
# ---------------------------------------------------------
send_status "CONFIGURING_RUNTIME" "Setting up NVIDIA Container Toolkit..."

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    sudo gpg --batch --yes --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' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

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

# ---------------------------------------------------------
# 4. Registry & DNS Configuration
# ---------------------------------------------------------
# Add registry.localhost to /etc/hosts if it doesn't already resolve to localhost
if ! getent hosts registry.localhost 2>/dev/null | grep -qE '(127\.0\.0\.1|::1)'; then
    echo "127.0.0.1 registry.localhost" | sudo tee -a /etc/hosts > /dev/null
fi

# Inject the private registry authentication
sudo mkdir -p /etc/rancher/k3s
cat <<EOF | sudo tee /etc/rancher/k3s/registries.yaml
mirrors:
  registry.localhost:
    endpoint:
      - https://registry.localhost:5000
configs:
  registry.localhost:
    auth:
      username: "$REG_USER"
      password: "$REG_PASS"
    tls:
      insecure_skip_verify: true
EOF

# ---------------------------------------------------------
# 5. Networking (Tailscale)
# ---------------------------------------------------------
send_status "CONNECTING_NETWORK" "Connecting to Tailscale..."
if ! command -v tailscale &> /dev/null; then
    curl -fsSL https://tailscale.com/install.sh | sudo sh
fi

sudo tailscale up --authkey="$TAILSCALE_KEY" --hostname="provider-$NODE_ID"
TS_IP=$(tailscale ip -4)

# ---------------------------------------------------------
# 6. K3s Agent Installation
# ---------------------------------------------------------
send_status "JOINING_CLUSTER" "Connecting node to GPU Orchestrator..."

curl -sfL https://get.k3s.io | K3S_URL="$K3S_URL" K3S_TOKEN="$K3S_TOKEN" sudo -E sh -s - \
  agent \
  --flannel-iface=tailscale0 \
  --node-ip="$TS_IP" \
  --node-name="provider-$NODE_ID"

# Final Cleanup
sudo rm -f "$PERSISTENT_SCRIPT"
sudo rm -f "$RESUME_FLAG"
send_status "SETUP_COMPLETE" "Setup successful!  Provider node is now online."
