install_aomp

#!/usr/bin/env bash
# Script to install the latest AOMP release for Ubuntu (amd64/x86_64)
# Usage: bash install_aomp [RETRY_MAX] [RETRY_DELAY_SECONDS]
#
# N.B.: Check https://github.com/ROCm/aomp/blob/aomp-dev/docs/MANYLINUX.md before running this
# script. Make sure that it is still the recommended installation method.

set -uo pipefail

# Default directory where AOMP is installed. Modify if needed.
AOMP_DIR="/usr/local/aomp"
echo -e "\n---------------------------------------------------------------------------------------------------"
echo -e "N.B.: AOMP will be installed at\n\n\t${AOMP_DIR}\n"
echo -e "After installation, please verify that AOMP is installed there by checking for the amdflang binary:\n"
echo -e "\tls ${AOMP_DIR}/bin/amdflang\n"
echo -e "If this directory is not desired, please modify the AOMP_DIR variable in the script and re-run it."
echo -e "---------------------------------------------------------------------------------------------------\n"


# The installation may fail due to network issues or GitHub API limits, so we retry until
# it works. We retry for at most $RETRY_MAX times with a delay of $RETRY_DELAY_SECONDS seconds
# between attempts.
if [[ "$#" -ge 1 ]]; then
    RETRY_MAX="$1"
else
    RETRY_MAX=5
fi
if [[ "$#" -ge 2 ]]; then
    RETRY_DELAY_SECONDS="$2"
else
    RETRY_DELAY_SECONDS=120
fi

# Allow update to fail (|| true) but proceed to install
sudo apt update || true
sudo apt install -y curl jq libelf-dev libzstd-dev  # jq needed for GitHub API parsing

# Make a temp directory and move into it
TMP_DIR="$(mktemp -d)"
cleanup() {
    # Only remove if it still exists
    [[ -n "${TMP_DIR:-}" && -d "${TMP_DIR}" ]] && rm -rf -- "${TMP_DIR}"
    TMP_DIR=""  # Prevent double cleanup
}

trap cleanup EXIT

echo "--> Created temporary directory: ${TMP_DIR}"
cd "${TMP_DIR}" || exit 1
echo "--> Changed to ${TMP_DIR}"

# Create a custom curl wrapper to enforce HTTP/1.1 and IPv4, with retries and timeouts.
# Useful in CI environments with occasional network issues.
cat >"$TMP_DIR/curl" <<'EOF'
#!/usr/bin/env bash
exec /usr/bin/curl --http1.1 -4 -L -C - --retry 5 --retry-connrefused --retry-delay 5 --no-keepalive --limit-rate 50M --connect-timeout 30 --max-time 1800 "$@"
EOF
chmod +x "$TMP_DIR/curl"
export PATH="$TMP_DIR:$PATH"
type curl

# Retrieve the latest AOMP release download URL from GitHub install page
INSTALL_PAGE="https://github.com/ROCm/aomp/blob/aomp-dev/docs/MANYLINUX.md"
URL_PATTERN="https://github.com/ROCm/aomp/releases.*?tar.gz"
DOWNLOAD_URL="$(curl -s ${INSTALL_PAGE} | grep -oP "${URL_PATTERN}" | sort -V | tail -n 1)"
if [[ -z "${DOWNLOAD_URL}" ]]; then
    echo "Error: Failed to retrieve the download URL from ${INSTALL_PAGE}" >&2
    exit 1
fi
FILENAME=$(basename "${DOWNLOAD_URL}")
echo "--> Retrieved info:"
echo "--> download URL: ${DOWNLOAD_URL}"
echo "--> filename: ${FILENAME}"

# Install AOMP
# See https://github.com/ROCm/aomp/blob/aomp-dev/docs/MANYLINUX.md
# It may fail due to network issues or GitHub API limits, so we retry until it works.
# We retry for at most $RETRY_MAX times with a delay of $RETRY_DELAY_SECONDS seconds between attempts.
RETRY_TIMES=0
AOMP_PARENT_DIR="$(dirname "${AOMP_DIR}")"

while true; do

    # Download the package
    curl -fL -o "${FILENAME}" "${DOWNLOAD_URL}"
    echo "--> Downloaded: ${FILENAME}"
    # Extract the package
    PACKAGE_NAME=$(tar -tf "${FILENAME}" | awk -F/ '{print $1}' | sort -u)
    tar -xzf "${FILENAME}"
    echo "--> Extracted: ${PACKAGE_NAME}"
    # Install the package
    sudo mv "${PACKAGE_NAME}" "$AOMP_PARENT_DIR/"
    sudo ln -sf "${AOMP_PARENT_DIR}/${PACKAGE_NAME}" "${AOMP_DIR}"
    # Verify installation
    if command -v $AOMP_DIR/bin/amdflang >/dev/null 2>&1 && $AOMP_DIR/bin/amdflang --version >/dev/null 2>&1; then
        echo "--> Installed AOMP to ${AOMP_DIR}"
        export PATH="$PATH:$AOMP_DIR/bin"
        break
    fi

    RETRY_TIMES=$((RETRY_TIMES + 1))

    if [ $RETRY_TIMES -ge "$RETRY_MAX" ]; then
        echo "Error: Installation failed after $RETRY_TIMES attempts." >&2
        exit 1
    else
        echo "amdflang does not work after installation attempt $RETRY_TIMES."
        echo "Retrying in $RETRY_DELAY_SECONDS seconds ... (Attempt: $RETRY_TIMES)"
        rm -f -- "${FILENAME}"
        sleep "$RETRY_DELAY_SECONDS"
    fi

done

# Final messages
echo "--> AOMP installation completed."
echo -e "\nThe amdflang version is:\n"
amdflang --version
echo -e "\nThe path to amdflang is:\n"
command -v amdflang
echo -e "\nTo use AOMP, you may need to add to your PATH:"
echo -e "\n\texport PATH=\$PATH:$AOMP_DIR/bin\n"
echo "Maybe also set LD_LIBRARY_PATH:"
echo -e "\n\texport LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:$AOMP_DIR/lib\n"
echo -e "Add these lines to your ~/.bashrc or ~/.profile for persistence.\n"

# Set environment variable for subsequent steps if running in GitHub Actions
if [[ -n "${GITHUB_ENV:-}" ]]; then
    echo "PATH=$PATH:$AOMP_DIR/bin" >> "$GITHUB_ENV"
fi

# Cleanup is handled by trap on EXIT