#!/bin/sh
# 9001-amd-rocm.hook.chroot — install AMD ROCm SMI tool for Instinct GPU monitoring.
# Runs inside the live-build chroot. Adds AMD's apt repository and installs
# rocm-smi-lib which provides the `rocm-smi` CLI (analogous to nvidia-smi).

set -e

# ROCm versions to try in order (newest first). Fall back if a version's
# Release file is missing from the repo (happens with brand-new releases).
ROCM_CANDIDATES="6.4 6.3 6.2"
ROCM_KEYRING="/etc/apt/keyrings/rocm.gpg"
ROCM_LIST="/etc/apt/sources.list.d/rocm.list"
APT_UPDATED=0

mkdir -p /etc/apt/keyrings

ensure_tool() {
    tool="$1"
    pkg="$2"
    if command -v "${tool}" >/dev/null 2>&1; then
        return 0
    fi
    if [ "${APT_UPDATED}" -eq 0 ]; then
        apt-get update -qq
        APT_UPDATED=1
    fi
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends "${pkg}"
}

ensure_cert_bundle() {
    if [ -s /etc/ssl/certs/ca-certificates.crt ]; then
        return 0
    fi
    if [ "${APT_UPDATED}" -eq 0 ]; then
        apt-get update -qq
        APT_UPDATED=1
    fi
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates
}

# live-build chroot may not include fetch/signing tools yet
if ! ensure_cert_bundle || ! ensure_tool wget wget || ! ensure_tool gpg gpg; then
    echo "WARN: failed to install wget/gpg/ca-certificates prerequisites — skipping ROCm install"
    exit 0
fi

# Download and import AMD GPG key
if ! wget -qO- "https://repo.radeon.com/rocm/rocm.gpg.key" \
        | gpg --dearmor --yes --output "${ROCM_KEYRING}"; then
    echo "WARN: failed to fetch AMD ROCm GPG key — skipping ROCm install"
    exit 0
fi

# Try each ROCm version until apt-get update succeeds (repo has a Release file).
ROCM_VERSION=""
for candidate in ${ROCM_CANDIDATES}; do
    cat > "${ROCM_LIST}" <<EOF
deb [arch=amd64 signed-by=${ROCM_KEYRING}] https://repo.radeon.com/rocm/apt/${candidate} bookworm main
EOF
    if apt-get update -qq 2>/dev/null; then
        ROCM_VERSION="${candidate}"
        echo "=== AMD ROCm ${ROCM_VERSION}: repository available ==="
        break
    fi
    echo "WARN: ROCm ${candidate} repository not available for bookworm, trying next..."
    rm -f "${ROCM_LIST}"
done

if [ -z "${ROCM_VERSION}" ]; then
    echo "WARN: no ROCm apt repository available for bookworm — skipping ROCm install"
    rm -f "${ROCM_KEYRING}"
    exit 0
fi

# rocm-smi-lib provides the rocm-smi CLI tool for GPU monitoring
if apt-get install -y --no-install-recommends rocm-smi-lib 2>/dev/null; then
    echo "=== AMD ROCm: rocm-smi installed ==="
    if [ -x /opt/rocm/bin/rocm-smi ]; then
        ln -sf /opt/rocm/bin/rocm-smi /usr/local/bin/rocm-smi
    else
        smi_path="$(find /opt -path '*/bin/rocm-smi' -type f 2>/dev/null | sort | tail -1)"
        if [ -n "${smi_path}" ]; then
            ln -sf "${smi_path}" /usr/local/bin/rocm-smi
        fi
    fi
    rocm-smi --version 2>/dev/null || true
else
    echo "WARN: rocm-smi-lib install failed — GPU monitoring unavailable"
fi

# Clean up apt lists to keep ISO size down
rm -f "${ROCM_LIST}"
apt-get clean