#!/bin/bash

set -u
set -e

# Check if machine is already attached to Ubuntu Pro
PRO_STATUS=$(pro status --format json 2>/dev/null | jq -r .attached 2>/dev/null)
if [ "${PRO_STATUS}" == "true" ]; then
    echo "Machine is already attached to Ubuntu Pro. No action needed."
    exit 0
fi

echo "Machine not attached to Ubuntu Pro. Attempting attachment with retries..."

# Create secure temporary files with restricted permissions
TOKEN_FILE1=$(mktemp)
TOKEN_FILE2=$(mktemp)
TOKEN_FILE3=$(mktemp)
chmod 600 "${TOKEN_FILE1}" "${TOKEN_FILE2}" "${TOKEN_FILE3}"

echo "Exporting Ubuntu Pro token..."
efivar --export "${TOKEN_FILE1}" --name="4f72e91a-a5b3-4c9d-8a6e-23d57bf4e9ac-UbuntuToken" || echo "Unable to export Ubuntu Pro token"
if [ -s "${TOKEN_FILE1}" ]; then
    echo "Extracting Ubuntu Pro token..."
    dd if="${TOKEN_FILE1}" of="${TOKEN_FILE2}" bs=1 skip=64 2>/dev/null
    dd if="${TOKEN_FILE2}" of="${TOKEN_FILE3}" bs=1 count=$(($(stat -c %s "${TOKEN_FILE2}") - 4)) 2>/dev/null
    TOKEN=$(cat "${TOKEN_FILE3}")

    echo "Waiting for network..."
    nm-online --wait-for-startup --timeout 300 || { echo "Network wait timed out"; exit 1; }

    # Retry attachment up to 10 times with 60 second delays
    for i in {1..10}; do
        echo "Attaching Ubuntu Pro token... (Attempt $i/10)"
        pro attach "${TOKEN}"
        echo "Ubuntu Pro token attached"

        # Check if attachment was successful
        PRO_STATUS=$(pro status --format json 2>/dev/null | jq -r .attached 2>/dev/null)
        if [ "${PRO_STATUS}" = "true" ]; then
            echo "Ubuntu Pro attachment successful!"
            break
        fi

        # If not the last attempt, wait 60 seconds
        if [ "$i" -lt 10 ]; then
            echo "Attachment may not have succeeded. Waiting 60 seconds before retry..."
            sleep 60
        else
            echo "Failed to attach Ubuntu Pro after 10 attempts"
            exit 1
        fi
    done
else
    echo "No Token present, disabling dgx-desktop-pro-activation.service"
    systemctl disable dgx-desktop-pro-activation.service || echo "Unable to disable dgx-desktop-pro-activation"
fi
