#!/bin/sh

# Configure script for ggmlR
# Vulkan backend: auto-detected, or forced via --with-vulkan / --without-vulkan
# SIMD flags: disabled by default, enabled via --with-simd

# Default values
USE_VULKAN="auto"
USE_SIMD="no"
VULKAN_CPPFLAGS=""
VULKAN_LIBS=""
VULKAN_OBJECTS=""

# Parse arguments
HOST=""
for arg in "$@"; do
  case "$arg" in
    --with-vulkan)
      USE_VULKAN="yes"
      ;;
    --without-vulkan)
      USE_VULKAN="no"
      ;;
    --with-simd)
      USE_SIMD="yes"
      ;;
    --host=*)
      HOST="${arg#--host=}"
      ;;
  esac
done

# --- Vulkan detection ---
#
# Three modes:
#   --with-vulkan    : require Vulkan, error if deps missing
#   --without-vulkan : disable Vulkan unconditionally
#   (default) auto   : enable if libvulkan-dev and glslc are found
#

if [ "$USE_VULKAN" = "no" ]; then
  echo "Vulkan: disabled by --without-vulkan"
fi

if [ "$USE_VULKAN" = "auto" ] || [ "$USE_VULKAN" = "yes" ]; then
  echo "Checking for Vulkan support..."

  # Check for glslc (shader compiler)
  # On Kaggle, prefer /usr/bin/glslc over /usr/local/bin/glslc (wrapper with bash syntax issues)
  if [ -x "/usr/bin/glslc" ]; then
    GLSLC="/usr/bin/glslc"
  else
    GLSLC=$(command -v glslc 2>/dev/null)
  fi

  # Check for Vulkan headers
  HAVE_VULKAN_H="no"
  if [ -f "/usr/include/vulkan/vulkan.h" ] || [ -n "$VULKAN_SDK" ]; then
    HAVE_VULKAN_H="yes"
  elif command -v pkg-config >/dev/null 2>&1 && pkg-config --exists vulkan 2>/dev/null; then
    HAVE_VULKAN_H="yes"
  fi

  if [ -z "$GLSLC" ] || [ "$HAVE_VULKAN_H" = "no" ]; then
    if [ "$USE_VULKAN" = "yes" ]; then
      # Explicit --with-vulkan: error out
      echo "ERROR: Vulkan dependencies not found."
      echo "  Required packages (Ubuntu/Debian): sudo apt install libvulkan-dev glslc"
      echo "  Or download Vulkan SDK from: https://vulkan.lunarg.com/sdk/home"
      [ -z "$GLSLC" ] && echo "  Missing: glslc (shader compiler)"
      [ "$HAVE_VULKAN_H" = "no" ] && echo "  Missing: vulkan/vulkan.h (Vulkan headers)"
      exit 1
    else
      # Auto mode: skip Vulkan with helpful message
      echo ""
      echo "Note: Vulkan GPU support not enabled (optional dependencies not found)."
      echo "  To enable Vulkan acceleration, install:"
      echo "    Ubuntu/Debian: sudo apt install libvulkan-dev glslc"
      echo "    Other systems: https://vulkan.lunarg.com/sdk/home"
      echo "  Building CPU-only version."
      echo ""
      USE_VULKAN="no"
    fi
  else
    USE_VULKAN="yes"
    if [ "$1" != "--with-vulkan" ]; then
      echo "Vulkan dev packages detected, enabling Vulkan automatically"
    fi
    echo "Found glslc: $GLSLC"
  fi
fi

# --- CPU SIMD flags ---
#
# CRAN policy forbids non-portable compilation flags like -mavx2, -msse4.2 etc.
# SIMD is disabled by default. Users who want SIMD acceleration on their own
# machine can enable it with:
#   install.packages("ggmlR", configure.args = "--with-simd")
#
SIMD_CFLAGS=""
if [ "$USE_SIMD" = "yes" ]; then
  echo "SIMD: enabled by --with-simd, detecting supported flags..."
  CC_CMD="${CC:-cc}"

  check_cflag() {
    if echo "int main(){return 0;}" | $CC_CMD $1 -x c -o /dev/null - 2>/dev/null; then
      SIMD_CFLAGS="$SIMD_CFLAGS $1"
      echo "  Found: $1"
      return 0
    fi
    return 1
  }

  SIMD_ARCH=$(uname -m 2>/dev/null || echo "unknown")
  case "$SIMD_ARCH" in
    x86_64|i386|i686|amd64)
      check_cflag -msse3
      check_cflag -mssse3
      check_cflag -msse4.1
      check_cflag -msse4.2
      check_cflag -mavx
      check_cflag -mavx2
      check_cflag -mfma
      check_cflag -mf16c
      ;;
    *)
      echo "  Non-x86 architecture ($SIMD_ARCH), skipping x86 SIMD flags"
      ;;
  esac

  if [ -z "$SIMD_CFLAGS" ]; then
    echo "  No SIMD flags detected, using scalar fallback"
  fi
else
  echo "SIMD: disabled (default). Use --with-simd to enable."
fi

# --- Detect OpenMP support ---
#
# We read SHLIB_OPENMP_CFLAGS and SHLIB_OPENMP_CXXFLAGS from R's Makeconf
# and substitute them as literal values into Makevars.in.
#
# Why not use SHLIB_OPENMP_* macros directly in Makevars?
# CRAN R CMD check enforces strict pairing rules for these macros.
# For mixed C/C++ packages with C++ linker, there is no valid combination:
#   - PKG_CFLAGS=$(SHLIB_OPENMP_CFLAGS) triggers "CFLAGS not in PKG_LIBS"
#   - PKG_CFLAGS=$(SHLIB_OPENMP_CXXFLAGS) triggers "incorrect macro in CFLAGS"
#   - Both macros in PKG_LIBS triggers "not portable to include multiple"
#
# By resolving flags at configure time (same approach as RcppArmadillo),
# the literal values (-fopenmp or empty) appear in Makevars, and R CMD check
# no longer sees the SHLIB_OPENMP_* macro names.
#
# On macOS (Apple clang): both values are empty -> no OpenMP, code uses
# pthreads fallback (#ifndef GGML_USE_OPENMP in ggml-cpu-backend.c).
#
echo "Checking OpenMP support..."
R_HOME="${R_HOME:-$(R RHOME)}"
MAKECONF="${R_HOME}/etc/Makeconf"

OPENMP_CFLAGS=""
OPENMP_CXXFLAGS=""
OPENMP_CPPFLAGS=""

# Emscripten (WASM) does not provide omp.h - disable OpenMP
if echo "$HOST" | grep -q "emscripten"; then
  echo "  Emscripten target detected, disabling OpenMP"
elif [ -f "$MAKECONF" ]; then
  OPENMP_CFLAGS=$(grep '^SHLIB_OPENMP_CFLAGS' "$MAKECONF" | sed 's/[^=]*= *//')
  OPENMP_CXXFLAGS=$(grep '^SHLIB_OPENMP_CXXFLAGS' "$MAKECONF" | sed 's/[^=]*= *//')
fi

if [ -n "$OPENMP_CFLAGS" ] || [ -n "$OPENMP_CXXFLAGS" ]; then
  OPENMP_CPPFLAGS="-DGGML_USE_OPENMP"
  echo "  OpenMP C flags: $OPENMP_CFLAGS"
  echo "  OpenMP C++ flags: $OPENMP_CXXFLAGS"
else
  echo "  OpenMP not available, using pthreads fallback"
fi

# --- Detect CPU architecture for arch-fallback ---
#
# We only ship arch/x86/ native implementations. On non-x86 platforms
# (ARM64 macOS, etc.) we need -DGGML_CPU_GENERIC so that arch-fallback.h
# maps all _generic functions to their canonical names, providing the
# symbols that would otherwise come from arch/aarch64/ or other arch dirs.
#
echo "Detecting CPU architecture..."
GENERIC_CPPFLAGS=""
X86_OBJECTS="ggml-cpu/arch/x86/cpu-feats.o ggml-cpu/arch/x86/quants.o ggml-cpu/arch/x86/repack.o"
ARCH=$(uname -m 2>/dev/null || echo "unknown")
case "$ARCH" in
  x86_64|i386|i686|amd64)
    echo "  x86 architecture detected, using native arch/x86/ implementations"
    ;;
  *)
    GENERIC_CPPFLAGS="-DGGML_CPU_GENERIC"
    X86_OBJECTS=""
    echo "  Non-x86 architecture ($ARCH), using generic CPU fallbacks"
    ;;
esac

# --- Build Vulkan shaders ---
if [ "$USE_VULKAN" = "yes" ]; then

  # Check for Vulkan headers in standard paths
  if [ ! -f "/usr/include/vulkan/vulkan.h" ] && [ -z "$VULKAN_SDK" ]; then
    echo "WARNING: vulkan/vulkan.h not found in standard paths."
    echo "If build fails, set VULKAN_SDK environment variable."
  fi

  # Check for pkg-config vulkan
  if command -v pkg-config >/dev/null 2>&1; then
    if pkg-config --exists vulkan 2>/dev/null; then
      VULKAN_CPPFLAGS=$(pkg-config --cflags vulkan)
      VULKAN_LIBS=$(pkg-config --libs vulkan)
      echo "Found Vulkan via pkg-config"
    fi
  fi

  # Fallback to manual detection
  if [ -z "$VULKAN_LIBS" ]; then
    VULKAN_LIBS="-lvulkan"
    if [ -n "$VULKAN_SDK" ]; then
      VULKAN_CPPFLAGS="-I${VULKAN_SDK}/include"
      VULKAN_LIBS="-L${VULKAN_SDK}/lib ${VULKAN_LIBS}"
    fi
  fi

  echo "Vulkan CPPFLAGS: $VULKAN_CPPFLAGS"
  echo "Vulkan LIBS: $VULKAN_LIBS"

  # Build shader generator
  echo "Building shader generator..."
  cd src/ggml-vulkan/vulkan-shaders

  CXX="${CXX:-g++}"
  $CXX -std=c++17 -O2 -o vulkan-shaders-gen vulkan-shaders-gen.cpp -pthread
  if [ $? -ne 0 ]; then
    echo "ERROR: Failed to build vulkan-shaders-gen"
    exit 1
  fi

  # Create output directories
  mkdir -p ../generated/spv

  # Generate header file
  echo "Generating shader header..."
  ./vulkan-shaders-gen \
    --output-dir ../generated/spv \
    --target-hpp ../generated/ggml-vulkan-shaders.hpp

  # Generate shader cpp files
  echo "Compiling shaders (this may take a while)..."
  SHADER_COUNT=0
  SHADER_TOTAL=$(ls -1 *.comp 2>/dev/null | wc -l)

  for shader in *.comp; do
    SHADER_COUNT=$((SHADER_COUNT + 1))
    printf "\r  [%d/%d] %s" "$SHADER_COUNT" "$SHADER_TOTAL" "$shader"
    ./vulkan-shaders-gen \
      --glslc "$GLSLC" \
      --source "$(pwd)/$shader" \
      --output-dir ../generated/spv \
      --target-hpp ../generated/ggml-vulkan-shaders.hpp \
      --target-cpp "../generated/${shader}.cpp"
    if [ $? -ne 0 ]; then
      echo ""
      echo "ERROR: Failed to compile shader: $shader"
      exit 1
    fi
  done
  echo ""
  echo "Generated $SHADER_COUNT shader files"

  cd ../../..

  # Build list of shader object files
  VULKAN_OBJECTS="ggml-vulkan/ggml-vulkan.o"
  for cpp in src/ggml-vulkan/generated/*.cpp; do
    base=$(basename "$cpp" .cpp)
    VULKAN_OBJECTS="$VULKAN_OBJECTS ggml-vulkan/generated/${base}.o"
  done

  VULKAN_CPPFLAGS="$VULKAN_CPPFLAGS -DGGML_USE_VULKAN -Iggml-vulkan -Iggml-vulkan/generated"

  echo "Vulkan backend enabled with $SHADER_COUNT shaders"
fi

# Generate Makevars from template
echo "Generating src/Makevars..."

sed \
  -e "s|@VULKAN_CPPFLAGS@|$VULKAN_CPPFLAGS|g" \
  -e "s|@VULKAN_LIBS@|$VULKAN_LIBS|g" \
  -e "s|@VULKAN_OBJECTS@|$VULKAN_OBJECTS|g" \
  -e "s|@OPENMP_CPPFLAGS@|$OPENMP_CPPFLAGS|g" \
  -e "s|@OPENMP_CFLAGS@|$OPENMP_CFLAGS|g" \
  -e "s|@OPENMP_CXXFLAGS@|$OPENMP_CXXFLAGS|g" \
  -e "s|@SIMD_CFLAGS@|$SIMD_CFLAGS|g" \
  -e "s|@GENERIC_CPPFLAGS@|$GENERIC_CPPFLAGS|g" \
  -e "s|@X86_OBJECTS@|$X86_OBJECTS|g" \
  src/Makevars.in > src/Makevars

echo "Configuration complete."
if [ "$USE_VULKAN" = "yes" ]; then
  echo "  Vulkan: ENABLED"
else
  echo "  Vulkan: disabled"
fi
if [ "$USE_SIMD" = "yes" ]; then
  echo "  SIMD:   ENABLED ($SIMD_CFLAGS )"
else
  echo "  SIMD:   disabled"
fi
