# NVIDIA open kernel module package feature.
#
# Public options:
#   _nvidia_pkg=""
#   _nvidia_sign=""

[[ -v _nvidia_pkg ]] || _nvidia_pkg=""
[[ -v _nvidia_sign ]] || _nvidia_sign=""

typeset -Ag _nvidia_versions=(
  [latest]="610.57.04"
  [vulkan]="595.44.09"
  [legacy]="580.173.02"
)
_nvidia_current_kernels=("6.18" "7.1" "7.2")

_resolve_nvidia() {
  _nvidia_var=""
  _nvidia_ver=""
  _nv_pkg=""
  local _variant

  case "${_nvidia_pkg:-}" in
    ""|"false")
      return 0
    ;;
    latest|vulkan|legacy)
      _nvidia_var="${_nvidia_pkg}"
      _nvidia_ver="${_nvidia_versions[${_nvidia_var}]}"
      _nvidia_pkg="${_nvidia_ver}"
      if [ "${_nvidia_var}" = "vulkan" ]; then
        _nv_pkg="open-gpu-kernel-modules-${_nvidia_ver}"
      else
        _nv_pkg="NVIDIA-kernel-module-source-${_nvidia_ver}"
      fi
      return 0
    ;;
  esac

  for _variant in latest vulkan legacy; do
    if [ "${_nvidia_pkg}" = "${_nvidia_versions[${_variant}]}" ]; then
      _nvidia_var="${_variant}"
      _nvidia_ver="${_nvidia_pkg}"
      if [ "${_nvidia_var}" = "vulkan" ]; then
        _nv_pkg="open-gpu-kernel-modules-${_nvidia_ver}"
      else
        _nv_pkg="NVIDIA-kernel-module-source-${_nvidia_ver}"
      fi
      return 0
    fi
  done

  _die "Invalid _nvidia_pkg value '${_nvidia_pkg}'. Use 'false' or one of: ${_nvidia_versions[latest]}, ${_nvidia_versions[vulkan]}, ${_nvidia_versions[legacy]}"
}

_set_nvidia_open() {
  if [ "${_nvidia_pkg}" = "false" ]; then
    _nvidia_var=""
    _nvidia_ver=""
    return
  fi

  if [ "${_numadisable}" = "true" ]; then
    warning "nvidia-open requires NUMA enabled for CUDA/NvEnc. _numadisable is \"true\" - disabling nvidia-open."
    _nvidia_pkg="false"
    _nvidia_var=""
    _nvidia_ver=""
    echo -e "_nvidia_pkg='false'\n_nvidia_var=''\n_nvidia_ver=''" >> "${_FROGMINER_TARGET:-${_where}/BIG_UGLY_FROGMINER}"
    return
  fi

  if [ -z "${_nvidia_pkg}" ]; then
      local _nvidia_prompt_array=()
      local _nvidia_var_map=()

      _nvidia_prompt_array+=("v${_nvidia_versions[latest]}")
      _nvidia_var_map+=("latest")
      _nvidia_prompt_array+=(
        "v${_nvidia_versions[vulkan]}"
        "v${_nvidia_versions[legacy]}"
        "skip building NVIDIA open kernel modules (default)"
      )
      _nvidia_var_map+=("vulkan" "legacy" "false")

      msg2 "Which open NVIDIA kernel module version do you want to build as a separate package?"
      msg2 "Requires Turing+ GPU (GTX 16xx and RTX 20xx or newer). Requires NUMA enabled."
      msg2 "Select an entry, or press Enter to skip."
      _default_index=$((${#_nvidia_prompt_array[@]} - 1))
      _prompt_from_array "${_nvidia_prompt_array[@]}"
      _nvidia_var="${_nvidia_var_map[$_selected_index]}"
      if [ "${_nvidia_var}" = "false" ]; then
        _nvidia_pkg="false"
      else
        _nvidia_pkg="${_nvidia_versions[${_nvidia_var}]}"
      fi
  fi

  if [ "${_nvidia_pkg}" = "false" ]; then
    _nvidia_var=""
    _nvidia_ver=""
    echo -e "_nvidia_pkg='false'\n_nvidia_var=''\n_nvidia_ver=''" >> "${_FROGMINER_TARGET:-${_where}/BIG_UGLY_FROGMINER}"
    return
  fi

  _resolve_nvidia

  if [ "${_nvidia_pkg}" != "false" ]; then
    local _kernel_supported="false"
    for _k in "${_nvidia_current_kernels[@]}"; do
      if [ "${_k}" = "${_basekernel}" ]; then
        _kernel_supported="true"
        break
      fi
    done
    if [ "${_kernel_supported}" != "true" ]; then
      error "nvidia-open is not supported for kernel ${_basekernel}. Supported kernels: ${_nvidia_current_kernels[*]}"
      error "Disabling nvidia-open."
      _nvidia_pkg="false"
      _nvidia_var=""
      _nvidia_ver=""
      echo -e "_nvidia_pkg='false'\n_nvidia_var=''\n_nvidia_ver=''" >> "${_FROGMINER_TARGET:-${_where}/BIG_UGLY_FROGMINER}"
      return
    fi

    if [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] && [ -z "${_nvidia_sign}" ]; then
      msg2 "[EXPERIMENTAL] Do you want to sign the NVIDIA open kernel modules?"
      msg2 "Requires CONFIG_MODULE_SIG=y in your kernel config. Has no effect without _nvidia_pkg."
      _nvidia_sign_array_text=("No (disable - recommended)" "Yes (enable - experimental)")
      _nvidia_sign_array=("false" "true")
      _default_index="0"
      _prompt_from_array "${_nvidia_sign_array_text[@]}"
      _nvidia_sign="${_nvidia_sign_array[$_selected_index]}"
      echo -e "_nvidia_sign='${_nvidia_sign}'" >> "${_FROGMINER_TARGET:-${_where}/BIG_UGLY_FROGMINER}"
    fi

    if [ "${_nvidia_pkg}" != "false" ]; then
      echo -e "_nvidia_pkg='${_nvidia_pkg}'\n_nvidia_var='${_nvidia_var}'\n_nvidia_ver='${_nvidia_ver}'" >> "${_FROGMINER_TARGET:-${_where}/BIG_UGLY_FROGMINER}"
    fi
  fi
}

_pkgname_extra() {
  [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] && printf '%s\n' "${pkgbase}-nvidia-open"
}

_source_extra() {
  [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] || return 0

  _resolve_nvidia
  if [ "${_nvidia_var}" = "vulkan" ]; then
    printf '%s.tar.gz::https://github.com/NVIDIA/open-gpu-kernel-modules/archive/refs/tags/%s.tar.gz\n' "${_nv_pkg}" "${_nvidia_ver}"
  else
    printf 'https://download.nvidia.com/XFree86/NVIDIA-kernel-module-source/%s.tar.xz\n' "${_nv_pkg}"
  fi
}

_build_after_kernel() {
  [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] || return 0

  _resolve_nvidia
  local _nv_open_src="${srcdir}/${_nv_pkg}"
  local _kernuname
  _kernuname="$(< "${_kernel_work_folder_abs}/include/config/kernel.release")"
  local MODULE_FLAGS=(
    KERNEL_UNAME="${_kernuname}"
    IGNORE_CC_MISMATCH=yes
    IGNORE_XEN_PRESENCE=1
    IGNORE_PREEMPT_RT_PRESENCE=1
    SYSSRC="${_kernel_work_folder_abs}"
    SYSOUT="${_kernel_work_folder_abs}"
  )
  if [ ! -d "${_nv_open_src}" ]; then
    _die "NVIDIA-open source directory not found: ${_nv_open_src}"
  fi
  msg2 "Building NVIDIA open kernel modules (${_nvidia_ver})..."
  env CFLAGS= CXXFLAGS= LDFLAGS= make "${BUILD_FLAGS[@]}" "${MODULE_FLAGS[@]}" \
    -C "${_nv_open_src}" -j"$(nproc)" modules
}

_nvidia_patch() {
  [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] || return 0

  _resolve_nvidia
  local _nv_open_src="${srcdir}/${_nv_pkg}"
  msg2 "NVIDIA-open-module source version ${_nvidia_ver} will be built and installed alongside this kernel."

  local _nv_patch="${srcdir}/0020-nvidia-${_nvidia_var}.patch"
  if [ ! -f "${_nv_patch}" ] && [ "${_nvidia_var}" = "vulkan" ]; then
    _nv_patch="${srcdir}/0020-nvidia-latest.patch"
  fi
  if [ -f "${_nv_patch}" ]; then
    msg2 "Applying NVIDIA-open-module patches (${_nvidia_ver})..."
    patch -Np1 -i "${_nv_patch}" -d "${_nv_open_src}"
  else
    msg2 "No nvidia patch found (looked for $(basename "${_nv_patch}"))."
  fi
}

_headers_should_cleanup_srcdir() {
  { [ "${_nvidia_pkg}" = "false" ] || [ -z "${_nvidia_pkg}" ]; } && { ! declare -F _module_split_enabled >/dev/null || ! _module_split_enabled; }
}

_package_nvidia_open() {
  source "${_where}/BIG_UGLY_FROGMINER"

  _resolve_nvidia

  pkgdesc="NVIDIA open modules driver $_nvidia_ver for the $pkgdesc kernel"
  depends=("${pkgbase}=${pkgver}" "libglvnd" "nvidia-utils-tkg>=${_nvidia_ver}")
  provides=("NVIDIA-MODULE")
  conflicts=("${pkgbase}-nvidia")
  license=('MIT AND GPL-2.0-only')

  _package_nvidia_open_files
}

_package_nvidia_open_files() {
  local _nv_open_src="${srcdir}/${_nv_pkg}"

  cd "${_kernel_work_folder_abs}"
  local _kernver="$(<version)"
  local modulesdir="${pkgdir}/usr/lib/modules/${_kernver}/extramodules"

  install -dm755 "${modulesdir}"
  install -m644 "${_nv_open_src}"/kernel-open/*.ko "${modulesdir}"
  install -Dt "${pkgdir}/usr/share/licenses/${pkgname}" -m644 "${_nv_open_src}/COPYING"

  local strip_bin="strip"
  [ "${_compiler_name}" = "-llvm" ] && strip_bin="llvm-strip"
  find "${modulesdir}" -type f -name '*.ko' -exec "${strip_bin}" --strip-debug '{}' \;

  if [[ "${_nvidia_sign}" == "true" ]]; then
    if declare -F _signing_resolve >/dev/null && _signing_resolve; then
      msg2 "Signing NVIDIA open kernel modules..."
      find "${modulesdir}" -type f -name '*.ko' -exec "${_sign_script}" "${_sign_hash}" "${_sign_key}" "${_sign_cert}" '{}' \;
    else
      warning "_nvidia_sign is enabled but signing is not available - skipping module signing."
    fi
  fi

  find "${pkgdir}" -name '*.ko' -exec zstd --rm -19 -T0 {} +

  if [ "${_NUKR}" = "true" ] && { ! declare -F _module_split_enabled >/dev/null || ! _module_split_enabled; }; then
    rm -rf "${srcdir}"
  fi
}

_package_functions() {
  [ "${_nvidia_pkg}" != "false" ] && [ -n "${_nvidia_pkg}" ] || return 0
  printf 'package_%s-nvidia-open() {\n  _package_nvidia_open\n}\n\n' "${pkgbase}"
}
