#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

echo "Exporting Ubuntu Pro token..."
/usr/bin/efivar --export /tmp/token1.txt --name="4f72e91a-a5b3-4c9d-8a6e-23d57bf4e9ac-UbuntuToken" || echo "Unable to export Ubuntu Pro token"
if [ -f "/tmp/token1.txt" ]; then
    echo "Extracting Ubuntu Pro token..."
    /usr/bin/dd if=/tmp/token1.txt of=/tmp/token2.txt bs=1 skip=64
    /usr/bin/dd if=/tmp/token2.txt of=/tmp/token3.txt bs=1 count=$(($(/usr/bin/stat -c %s /tmp/token2.txt) - 4))
    TOKEN=$(/usr/bin/cat /tmp/token3.txt)

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

    # Check if machine is already attached to Ubuntu Pro
    PRO_STATUS=$(/usr/bin/pro status --format json 2>&1 | /usr/bin/jq .attached)

    if [ "${PRO_STATUS}" != "true" ]; then
        echo "Machine not attached to Ubuntu Pro. Attempting attachment with retries..."

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

            # Check if attachment was successful
            PRO_STATUS=$(/usr/bin/pro status --format json 2>&1 | /usr/bin/jq .attached)
            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 0
            fi
        done
    else
        echo "Machine is already attached to Ubuntu Pro or pro command not available"
    fi
    rm -f /tmp/token?.txt
    /usr/bin/systemctl disable dgx-oobe-ubuntu-pro-activate.service || echo "Unable to disable dgx-oobe-ubuntu-pro-activate"
else
    echo "No Token present, disabling dgx-oobe-ubuntu-pro-activate.service"
    /usr/bin/systemctl disable dgx-oobe-ubuntu-pro-activate.service || echo "Unable to disable dgx-oobe-ubuntu-pro-activate"
fi
