#!/bin/sh

# Configure CMCMC for Unix-like platforms.
#
# If a CUDA toolkit is found, build the real CUDA backend. Otherwise build a
# small stub backend so the R package can still be installed and checked on
# machines without CUDA. The C CPU backend for cmcmc and inca is always built.
# OpenMP is optional. If it is not available, the CPU backend is built in
# serial mode and the nthreads argument is ignored.

set -eu

cuda_arg=""

for arg in "$@"; do
  case "$arg" in
    --with-cuda=*)
      cuda_arg=${arg#--with-cuda=}
      ;;
  esac
done

first_existing_dir() {
  for d in "$@"; do
    if test -n "$d" && test -d "$d"; then
      printf '%s\n' "$d"
      return 0
    fi
  done
  return 1
}

r_cmd_config() {
  if test -n "${R_HOME:-}" && test -x "${R_HOME}/bin/R"; then
    "${R_HOME}/bin/R" CMD config "$1" 2>/dev/null || true
  else
    R CMD config "$1" 2>/dev/null || true
  fi
}

find_libomp_home() {
  first_existing_dir \
    "${LIBOMP_HOME:-}" \
    "${HOMEBREW_PREFIX:-}/opt/libomp" \
    /opt/homebrew/opt/libomp \
    /usr/local/opt/libomp
}

detect_openmp() {
  OPENMP_CFLAGS=""
  OPENMP_LIBS=""

  case "${CMCMC_OPENMP:-auto}" in
    no|NO|false|FALSE|0)
      echo "OpenMP disabled via CMCMC_OPENMP=${CMCMC_OPENMP}; building serial CPU backend."
      return 0
      ;;
  esac

  if test -n "${CMCMC_OPENMP_CFLAGS:-}" || test -n "${CMCMC_OPENMP_LIBS:-}"; then
    OPENMP_CFLAGS="${CMCMC_OPENMP_CFLAGS:-}"
    OPENMP_LIBS="${CMCMC_OPENMP_LIBS:-$OPENMP_CFLAGS}"
    echo "OpenMP flags set from CMCMC_OPENMP_CFLAGS/CMCMC_OPENMP_LIBS."
    return 0
  fi

  OPENMP_CFLAGS=$(r_cmd_config SHLIB_OPENMP_CFLAGS)
  case "$OPENMP_CFLAGS" in
    ""|ERROR:*)
      OPENMP_CFLAGS=""
      ;;
  esac
  if test -n "$OPENMP_CFLAGS"; then
    OPENMP_LIBS="$OPENMP_CFLAGS"
    echo "OpenMP detected from R configuration: $OPENMP_CFLAGS"
    return 0
  fi

  case "$(uname -s)" in
    Darwin)
      libomp_home=$(find_libomp_home || true)
      if test -n "$libomp_home" &&
         test -f "$libomp_home/include/omp.h" &&
         { test -f "$libomp_home/lib/libomp.dylib" || test -f "$libomp_home/lib/libomp.a"; }; then
        OPENMP_CFLAGS="-Xpreprocessor -fopenmp -isystem $libomp_home/include"
        OPENMP_LIBS="-L$libomp_home/lib -lomp"
        echo "OpenMP detected via Homebrew libomp at $libomp_home"
        return 0
      fi
      ;;
  esac

  echo "OpenMP not detected; building serial CPU backend."
}

find_cuda_home() {
  if test -n "$cuda_arg" && test "$cuda_arg" != "no"; then
    first_existing_dir "$cuda_arg"
    return $?
  fi

  if test -n "${CUDA_HOME:-}"; then
    first_existing_dir "$CUDA_HOME" && return 0
  fi
  if test -n "${CUDA_PATH:-}"; then
    first_existing_dir "$CUDA_PATH" && return 0
  fi
  if test -n "${CUDA_ROOT:-}"; then
    first_existing_dir "$CUDA_ROOT" && return 0
  fi

  first_existing_dir \
    /usr/local/cuda \
    /usr/local/cuda-* \
    /opt/cuda \
    /Developer/NVIDIA/CUDA-*
}

write_cuda_makevars() {
  {
    printf 'CUDA_HOME = %s\n' "$CUDA_HOME"
    printf 'CUDA_LIB_LIB64 = %s\n' "$CUDA_LIB_LIB64"
    printf 'CMCMC_OPENMP_CFLAGS = %s\n' "$OPENMP_CFLAGS"
    printf 'CMCMC_OPENMP_LIBS = %s\n' "$OPENMP_LIBS"
    cat <<'EOF'
OBJECTS = cmcmcR.o incaR.o init.o cmcmc_cpu.o
PKG_LIBS = -L./cmcmc -lcmcmc -L$(CUDA_HOME)/$(CUDA_LIB_LIB64) -Wl,-rpath,"$(CUDA_HOME)/$(CUDA_LIB_LIB64)" -lcudart -lcurand -lcublas -lcusolver $(CMCMC_OPENMP_LIBS)
PKG_CFLAGS = -I$(CUDA_HOME)/include -DADD_ $(CMCMC_OPENMP_CFLAGS)

.PHONY: all

all: $(SHLIB)
$(SHLIB): cmcmc/libcmcmc.a

cmcmc/libcmcmc.a:
	@(cd cmcmc && $(MAKE) libcmcmc.a \
		CC="$(CC)" CFLAGS="$(CFLAGS) $(CPICFLAGS) -I$(R_INCLUDE_DIR) -DRUN_IN_R" \
		AR="$(AR)" RANLIB="$(RANLIB)" CUDA_HOME="$(CUDA_HOME)" \
		)
EOF
  } > src/Makevars
}

write_stub_makevars() {
  {
    printf 'CMCMC_OPENMP_CFLAGS = %s\n' "$OPENMP_CFLAGS"
    printf 'CMCMC_OPENMP_LIBS = %s\n' "$OPENMP_LIBS"
    cat <<'EOF'
OBJECTS = cmcmcR.o incaR.o init.o cmcmc_stub.o cmcmc_cpu.o
PKG_CFLAGS = -DADD_ $(CMCMC_OPENMP_CFLAGS)
PKG_LIBS = $(CMCMC_OPENMP_LIBS)
EOF
  } > src/Makevars
}

detect_openmp

if test "$cuda_arg" = "no"; then
  echo "CUDA disabled via --with-cuda=no; building CMCMC stub backend."
  write_stub_makevars
  exit 0
fi

CUDA_HOME=$(find_cuda_home || true)
CUDA_LIB_LIB64=""

if test -n "$CUDA_HOME" && test -d "$CUDA_HOME/include"; then
  if test -d "$CUDA_HOME/lib64"; then
    CUDA_LIB_LIB64=lib64
  elif test -d "$CUDA_HOME/lib"; then
    CUDA_LIB_LIB64=lib
  fi
fi

if test -n "$CUDA_HOME" && test -n "$CUDA_LIB_LIB64"; then
  echo "CUDA detected at $CUDA_HOME; building CMCMC CUDA backend."
  write_cuda_makevars
else
  echo "CUDA toolkit not found; building CMCMC stub backend."
  write_stub_makevars
fi

exit 0
