#!/bin/bash
#
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

# Print usage information, then quit
usage() {
  echo "Usage: $0 [bindir] [switches]..."
  echo "bindir       the bin directory where the NVIDIA HPC SDK compilers are installed"
  echo "-d outdir    directory to put localrc; default is bindir"
  echo "-n           show the install information, do not write a localrc file"
  echo "-o           create the localrc file and write to stdout"
  echo "-x           write the localrc file and place in outputdir"
  echo "-q           quiet mode"
  echo "-v           verbose mode"
  echo "-gcc gcc     gcc command name; default is gcc"
  echo "-gpp g++     g++ command name; default is g++"
  echo "-g77 g77     g77 command name; default is g77"
  echo "-net locdir  perform network install with named local directory"
  echo "-cuda ver    set default cuda version, e.g. 11.0, 11.8, 12.3"
  echo "Example: $0 -x $(dirname $(which $0))"
  exit 1
}

umask 022

LANG=C
unset CPATH
unset C_INCLUDE_PATH
unset CPLUS_INCLUDE_PATH

# defaults
defcudaversion=""
echoout=false
gcc="gcc"
gpp="g++"
gfortran="gfortran"
ldd="ldd"
localrc="localrc"
newrc=$(mktemp)
network=false
noexec=true
nvaccelinfo="nvaccelinfo"
pathregex='[[:alnum:]/_.-]+'
compat=true
cuda=true
sanity=true

# cleanup temporary file
trap "rm -f ${newrc}" EXIT

# basic Fortran program
IFS='' read -r -d '' HELLO_F <<EOF
      program hello
      write(*,100)
100   format('Hello world!')
      end
EOF

# basic C program
IFS='' read -r -d '' NULL_C <<EOF
int main() {}
EOF

error() {
  # print error and quit
  local msg=$*
  echo -e "ERROR: ${msg}" 1>&2
  exit 1
}

# return the path to a library given a link flag (-l<lib>)
find_lib() {
  local lib=$1
  echo "${NULL_C}" | ${gcc} ${lib} -x c -o /dev/null - >& /dev/null
  if [ $? -eq 0 ]; then
    echo "${lib}"
  else
    echo ""
  fi
}

# return the include path for a specified language
find_inc() {
  local lang=$1
  cpp_v=$(${gcc} -E -Wp,-v -x ${lang} /dev/null 2>&1)
  regex='starts.here:.[^#](.*).End\ '
  [[ ${cpp_v} =~ ${regex} ]] && echo "${BASH_REMATCH[1]}"
}

# print a line to the localrc file and/or stdout
print_line() {
  local msg=$*

  if ! $noexec; then
    echo ${msg} >> $newrc
  fi

  if $echoout; then
    echo ${msg}
  fi
}

# verify that all dependencies are present
sanity_check() {
  if ! type flock &> /dev/null; then
    error "flock not found"
  fi

  if ! type find &> /dev/null; then
    error "find not found"
  fi

  if ! type "${gcc}" &> /dev/null; then
    error "gcc not found"
  fi

  if ! type "${gpp}" &> /dev/null; then
    error "g++ not found"
  fi

  if ! type install &> /dev/null; then
    error "install not found"
  fi
  
  if ! type "${ldd}" &> /dev/null; then
    error "ldd not found"
  fi

  if ! type mktemp &> /dev/null; then
    error "mktemp not found"
  fi
}

# set the CUDA parameter DEFCUDAVERSION
set_cuda() {
  local cudaversion="$1"

  # DEFCUDAVERSION
  if [ -n "${cudaversion}" ]; then
    print_line "set DEFCUDAVERSION=${cudaversion};"
  else
    # check for nvaccelinfo
    if ! type "${bindir}/nvaccelinfo" &> /dev/null; then
      print_line "# nvaccelinfo not found"
      print_line "#set DEFCUDAVERSION=;"
      return
    fi

    # query the GPU
    nvaccelinfo_out=$($bindir/nvaccelinfo -dev 0 -mig)
    nvaccelinfo_rt=$?

    [[ ${nvaccelinfo_out} =~ CUDA\ Driver\ Version:\ +([[:digit:]]+) ]] && CUDA_DRIVER_VERSION="${BASH_REMATCH[1]}"

    if ! test -d ${bindir}/../../cuda ; then
      error "cuda directory not found"
    fi

    # see what version(s) of CUDA are bundled with this install
    readarray -t bundled_cuda < <(find $bindir/../../cuda -maxdepth 1 -type d -regex '.*/[0-9]+\.[0-9]+' -printf '%f\n' | sort -n)

    if [ -z "$nvaccelinfo_out" ] || [ "${nvaccelinfo_rt}" -ne 0 ]; then
      # No GPUs present
      DEFCUDAVERSION=${bundled_cuda[-1]}
    else
      # Use the maximum bundled CUDA version that does not exceed the driver version
      for cudaver in "${bundled_cuda[@]}"; do
        if [ "${CUDA_DRIVER_VERSION}" -ge "${cudaver//./0}0" ]; then
          DEFCUDAVERSION="${cudaver}"
        else
          # Handle case where minimum bundled CUDA version exceeds the driver version
          if [ -z "$DEFCUDAVERSION" ]; then
            DEFCUDAVERSION=${bundled_cuda[-1]}
          fi
          break
        fi
      done
    fi

    print_line "set DEFCUDAVERSION=${DEFCUDAVERSION};"
  fi
}

###
# parse command line
###
while test -n "$1"; do # {
  case "$1" in
  -h )   usage ;;
  -n )   echoout=true ; noexec=true ;;
  -x )   noexec=false ;;
  -q )   quiet=1 ; verbose=0 ;;
  -v )   verbose=1 ; quiet=0 ;;
  -o )   echoout=true ; noexec=false ;;
  -d )   shift; outputdir="$1" ;;
  -f )   shift; localrc="$1" ;;
  -net ) shift; locdir="$1"; network=true ;;
  -lnet) locdir= ; network=true ;; # undocumented, internal testing only
  -gcc ) shift; gcc="$1" ;;
  -gpp ) shift; gpp="$1" ;;
  -g77 ) shift; gfortran="$1" ;;
  -cuda ) shift; defcudaversion="$1" ;;
  -stdpar ) shift; error "-stdpar is no longer supported" ;;
  -no-cuda ) cuda=false ;; # undocumented, internal testing only 
  -skip-sanity ) sanity=false ;; # undocumented, internal testing only
  -drop-compat ) compat=false ;; # undocumented, internal testing only
  * )    bindir="$1" ;;
  esac
  shift
done

if $noexec && ! $echoout; then
  usage
fi

if [ -z "${bindir}" ]; then 
  bindir=$(dirname $0)
fi

if [ ! -d "${bindir}" ]; then
  error "${bindir}: directory not found"
fi

# check that all dependences are present
if $sanity; then
  sanity_check
fi

# LFC
if type ${gfortran} &> /dev/null; then
  gfortran_hello=$(echo "${HELLO_F}" | ${gfortran} -v -x f77 -o /dev/null - 2>&1) || error "${gfortran_hello}"
  case ${gfortran_hello} in
    *"lgfortran"* | *"l gfortran"*)
      LFC="-lgfortran"
      ;;
    *"lf2c"*)
      LFC="-lf2c"
      ;;
    *"lg2c"*)
      LFC="-lg2c"
      ;;
    *)
      LFC=""
  esac
  print_line "set LFC=${LFC};"
fi

# LDSO
gcc_v=$(echo "${NULL_C}" | ${gcc} -v -x c -o /dev/null - 2>&1) || error "${gcc_v}"
[[ ${gcc_v} =~ -dynamic-linker\ +($pathregex) ]] && LDSO=${BASH_REMATCH[1]}
print_line "set LDSO=${LDSO};"

# GCCDIR
gcc_print_search_dirs=$(${gcc} -print-search-dirs) || error "$gcc_print_search_dirs"
[[ ${gcc_print_search_dirs} =~ install:\ +($pathregex) ]] && GCCDIR=${BASH_REMATCH[1]}
print_line "set GCCDIR=${GCCDIR};"

# G77DIR
if type ${gfortran} &> /dev/null; then
  gfortran_print_search_dirs=$(${gfortran} -print-search-dirs) || error "${gfortran_print_search_dirs}"
  [[ ${gfortran_print_search_dirs} =~ install:\ +($pathregex) ]] && GFORTRANDIR=${BASH_REMATCH[1]}
  print_line "set G77DIR=${GFORTRANDIR};"
fi

# OEM_INFO
arch=$(uname -s -m)
arch=${arch// /_} # replace space with underscore
case "$arch" in
  Linux_x86_64 )
    OEM_INFO="64-bit target on x86-64 Linux \$INFOTPVAL"
    ;;
  Linux_ppc64le )
    OEM_INFO="linuxpower target on Linuxpower \$INFOTPVAL"
    ;;
  Linux_aarch64 )
    OEM_INFO="linuxarm64 target on aarch64 Linux \$INFOTPVAL"
    ;;
  * )
    error "unexpected architecture '${arch}'"
    ;;
esac
print_line "set OEM_INFO=$OEM_INFO;"

# GNUATOMIC
GNUATOMIC=$(find_lib -latomic)
print_line "set GNUATOMIC=${GNUATOMIC};"

# GCCINC
GCCINC=$(find_inc c)
print_line "set GCCINC=${GCCINC};"

# GPPDIR
GPPDIR=$(find_inc c++)
print_line "set GPPDIR=${GPPDIR};"

# NUMALIBNAME
if $compat; then
  if [ "${arch}" = "Linux_x86_64" ]; then
    NUMALIBNAME=$(find_lib -lnuma)
    print_line "set NUMALIBNAME=${NUMALIBNAME};"
  fi
fi

# LOCALRC
print_line "set LOCALRC=YES;"

if $compat; then
  ldd_version=$(${ldd} --version 2> /dev/null) || error "${ldd_version}"
  [[ ${ldd_version} =~ ^ldd\ +\(.*\)\ +([0-9\.]+)$'\n' ]] && GLIBC_VERSION="${BASH_REMATCH[1]}"

  case "$GLIBC_VERSION" in
    2.[3-9]* | 2.[1-2][0-9]* )
      glibc=232
      ;;
    0.* | 1.* | 2.0.* | 2.1.* | 2.2.*)
      error "glibc version '$GLIBC_VERSION' is no longer supported."
      ;;
    2.* )
      echo "Unknown glibc version '$GLIBC_VERSION'; treating like 2.22" 1>&2
      glibc=222
      ;;
    * )
      error "unknown glibc version '$GLIBC_VERSION'"
      ;;
  esac
else
  # Everything is a modern enough GLIBC now, right?
  glibc="yes"
fi

if [ -n "$glibc" ]; then
  # EXTENSION
  print_line "set EXTENSION=__extension__=;"

  # COMPGLIBLFDIR
  if $compat; then
    if [ "$arch" != "Linux_x86_64" ]; then
      print_line "set COMPGLIBLFDIR=;"
    fi
  fi

  # LC
  if [ -e "$GCCDIR/libgcc_eh.a" ]; then
    case "$arch" in
      Linux_x86_64 )
        print_line "set LC=-lgcc -lc \$if(-Bstatic,-lgcc_eh, -lgcc_s);"
        ;;
      Linux_aarch64 | Linux_ppc64le )
        print_line "set LC=\$if(-Bstatic,-lgcc -lgcc_eh -lc, -lgcc --as-needed -lgcc_s --no-as-needed -lc -lgcc --as-needed -lgcc_s --no-as-needed);"
        ;;
    esac
  else
    case "$arch" in
      Linux_x86_64 )
        print_line "set LC=-lgcc -lc -lgcc_s;"
        ;;
      Linux_aarch64 | Linux_ppc64le )
        print_line "set LC=-lgcc --as-needed -lgcc_s --no-as-needed -lc -lgcc --as-needed -lgcc_s --no-as-needed;"
        ;;
    esac
  fi

  case "${arch}" in 
    Linux_x86_64 )
      libdirs="/usr/lib/x86_64-linux-gnu"
      ;;
    Linux_aarch64 )
      DEFLIBDIR="/usr/lib"
      libdirs="/usr/lib/aarch64-linux-gnu /usr/lib/arm-linux-gnueabi /usr/lib64"
      ;;
    Linux_ppc64le )
      DEFLIBDIR="/usr/lib"
      libdirs="/usr/lib/powerpc64le-linux-gnu /usr/lib/power-linux-gnueabi /usr/lib64"
      ;;
  esac

  for libdir in ${libdirs}; do
    if [ -e "${libdir}/crt1.o" ]; then
      DEFLIBDIR="${libdir}"
      break
    fi
  done

  # DEFLIBDIR
  if [ "${arch}" != "Linux_x86_64" ]; then
    print_line "set DEFLIBDIR=${DEFLIBDIR};"
  fi

  # DEFSTDOBJDIR
  if [ -n "$DEFLIBDIR" ]; then
    print_line "set DEFSTDOBJDIR=${DEFLIBDIR};"
  fi
fi

if $cuda; then
  set_cuda "$defcudaversion"
fi

print_line "# GLIBC version ${GLIBC_VERSION}"

# GCCVERSION
gcc_version=$(${gcc} -dumpfullversion 2>&1)
if [ $? -ne 0 ]; then
  gcc_version=$(${gcc} -dumpversion 2>&1)
fi
GCCVERSION=${gcc_version//./0}
print_line "# GCC version ${gcc_version}"
print_line "set GCCVERSION=${GCCVERSION};"

# GCCNAME
GCCNAME="${gcc##*/}" # basename
if [ "${GCCNAME}" != "gcc" ]; then
  print_line "set GCCNAME=${GCCNAME};"
fi

# LOCALDEFS
if [ "${arch}" != "Linux_x86_64" ]; then
  if [ "$GCCVERSION" -gt "40500" ]; then
    print_line "set LOCALDEFS=__STDC_HOSTED__;"
  fi
fi

# PGI
print_line "export PGI=\$COMPBASE;"

# footer
print_line "# makelocalrc executed by ${USER} $(date)"

# all done, install localrc
if ! $noexec; then
  if $network; then
    hostname=$(hostname -s)
    localrc="$localrc.$hostname" 
  fi

  if [ -n "$outputdir" ]; then
    mkdir -p "$outputdir"
    if [ ! -d "$outputdir" ]; then
      error "output directory '$outputdir' not found"
    fi
    localrc="$outputdir/$localrc"
  else
    if [ ! -d "$bindir" ]; then
      error "output directory '$bindir' not found"
    fi
    localrc="$bindir/$localrc"
  fi

  flock "${localrc}-lock" install --mode=644 --backup --suffix=.bak "$newrc" "$localrc"
  if [ $? -ne 0 ]; then
    error "unable to create ${localrc}"
  fi
fi
