hwaccel.func

#!/usr/bin/env bash
# Copyright (c) 2021-2026 community-scripts ORG
# License: MIT | https://raw.githubusercontent.com/community-scripts/core/main/LICENSE

# ==============================================================================
# TOOLS/HWACCEL.FUNC - GPU DETECTION AND HARDWARE ACCELERATION
# ==============================================================================
# Part of the shared helper library. Load lib/tools.func -- it sources every
# part -- rather than this file directly; the split is for maintenance only and
# no function name or behaviour changed with it.
# ==============================================================================

[[ -n "${_TOOLS_HWACCEL_LOADED:-}" ]] && return 0
_TOOLS_HWACCEL_LOADED=1

# Install a best-effort package list without letting one absentee take the rest.
# apt is all-or-nothing, so one package missing on a release meant none of the
# other eleven were installed and the container ran without acceleration on a
# single warning line. The batch is still tried first -- one apt run, and the
# normal case -- and only a failure falls back to one at a time, which is also
# the only way to name the package that was actually missing.
_cs_apt_install_optional() {
  (($#)) || return 0
  $STD apt -y install "$@" 2>/dev/null && return 0

  local missed=() failed=() pkg
  for pkg in "$@"; do
    $STD apt -y install "$pkg" 2>/dev/null && continue
    # A broken apt state fails every install, which is not the same thing.
    if apt-cache policy "$pkg" 2>/dev/null | grep -q "Candidate: (none)"; then
      missed+=("$pkg")
    else
      failed+=("$pkg")
    fi
  done
  ((${#missed[@]})) && msg_warn "Not available on this release: ${missed[*]}"
  ((${#failed[@]})) && msg_warn "In the repository but would not install: ${failed[*]} - apt may be in a broken state"
  return 0
}

# ------------------------------------------------------------------------------
# Sets up Hardware Acceleration on debian or ubuntu.
#
# Description:
#   - Detects all available GPUs (Intel, AMD, NVIDIA)
#   - Allows user to select which GPU(s) to configure (with 60s timeout)
#   - Installs the correct libraries and packages for each GPU type
#   - Supports: Debian 11/12/13, Ubuntu 22.04/24.04
#   - Intel: Legacy (Gen 6-8), Modern (Gen 9+), Arc
#   - AMD: Discrete GPUs, APUs, ROCm compute
#   - NVIDIA: Version-matched drivers from CUDA repository
#
# Notes:
#   - Some Intel packages are fetched from GitHub due to missing Debian packages
#   - NVIDIA requires matching host driver version
# ------------------------------------------------------------------------------
setup_hwaccel() {
  local service_user="${1:-}"

  export HWACCEL_VENDOR="none"
  _cs_nvidia_usable=0

  # Check if user explicitly disabled GPU in advanced settings
  # ENABLE_GPU is exported from build.func
  if [[ "${ENABLE_GPU:-no}" == "no" ]]; then
    return 0
  fi

  # Check if GPU passthrough is enabled (device nodes must exist)
  if [[ ! -d /dev/dri && ! -e /dev/nvidia0 && ! -e /dev/kfd ]]; then
    msg_warn "No GPU passthrough detected (/dev/dri, /dev/nvidia*, /dev/kfd not found) - skipping hardware acceleration setup"
    return 0
  fi

  msg_info "Setup Hardware Acceleration"

  # Install pciutils if needed
  if ! command -v lspci &>/dev/null; then
    $STD apt -y update || {
      msg_warn "Failed to update package list"
      return 0
    }
    $STD apt -y install pciutils || {
      msg_warn "Failed to install pciutils"
      return 0
    }
  fi

  # ═══════════════════════════════════════════════════════════════════════════
  # GPU Detection - Build list of all available GPUs with details
  # ═══════════════════════════════════════════════════════════════════════════
  local -a GPU_LIST=()
  local -a GPU_TYPES=()
  local -a GPU_NAMES=()
  local gpu_count=0

  # Get all GPU entries from lspci
  while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    local pci_addr gpu_name gpu_type=""

    pci_addr=$(echo "$line" | awk '{print $1}')
    gpu_name=$(echo "$line" | sed 's/^[^ ]* [^:]*: //')

    # Determine GPU type
    # Note: Use -w (word boundary) for ATI to avoid matching "CorporATIon"
    if echo "$gpu_name" | grep -qi 'Intel'; then
      gpu_type="INTEL"
      # Subtype detection for Intel
      # Order matters: Check Arc first, then Gen9+ (UHD/Iris/HD 5xx-6xx), then Legacy (HD 2xxx-5xxx)
      # HD Graphics 530/630 = Gen 9 (Skylake/Kaby Lake) - 3 digits
      # HD Graphics 4600/5500 = Gen 7-8 (Haswell/Broadwell) - 4 digits starting with 2-5
      if echo "$gpu_name" | grep -qiE 'Arc|DG[12]'; then
        gpu_type="INTEL_ARC"
      elif echo "$gpu_name" | grep -qiE 'UHD|Iris|HD Graphics [5-6][0-9]{2}[^0-9]|HD Graphics [5-6][0-9]{2}$'; then
        # HD Graphics 5xx/6xx (3 digits) = Gen 9+ (Skylake onwards)
        gpu_type="INTEL_GEN9+"
      elif echo "$gpu_name" | grep -qiE 'HD Graphics [2-5][0-9]{3}'; then
        # HD Graphics 2xxx-5xxx (4 digits) = Gen 6-8 Legacy
        gpu_type="INTEL_LEGACY"
      fi
    elif echo "$gpu_name" | grep -qiwE 'AMD|ATI|Radeon|Advanced Micro Devices'; then
      gpu_type="AMD"
    elif echo "$gpu_name" | grep -qi 'NVIDIA'; then
      # lspci shows the host's bus; the card is only ours with its device nodes.
      if compgen -G '/dev/nvidia[0-9]*' >/dev/null; then
        gpu_type="NVIDIA"
      else
        msg_custom "ℹ️" "${YW}" "NVIDIA GPU on the host, but no /dev/nvidia* device in this container - skipping it (add the devices to the container config to use it)"
      fi
    fi

    if [[ -n "$gpu_type" ]]; then
      GPU_LIST+=("$pci_addr")
      GPU_TYPES+=("$gpu_type")
      GPU_NAMES+=("$gpu_name")
      ((gpu_count++)) || true
    fi
  # The class field only: the whole line matched "Device 43d0" and "SanDisk Ultra 3D".
  done < <(lspci 2>/dev/null | grep -E '^[^ ]+ (VGA compatible|XGA compatible|3D|Display) controller[ :]')

  # Check for AMD APU via CPU vendor if no discrete GPU found
  local cpu_vendor
  cpu_vendor=$(lscpu 2>/dev/null | grep -i 'Vendor ID' | awk '{print $3}' 2>/dev/null || echo "")

  if [[ $gpu_count -eq 0 ]]; then
    if [[ "$cpu_vendor" == "AuthenticAMD" ]]; then
      GPU_LIST+=("integrated")
      GPU_TYPES+=("AMD_APU")
      GPU_NAMES+=("AMD APU (Integrated Graphics)")
      ((gpu_count++)) || true
    else
      msg_warn "No GPU detected - skipping hardware acceleration setup"
      return 0
    fi
  fi

  # ═══════════════════════════════════════════════════════════════════════════
  # GPU Selection - Let user choose which GPU(s) to configure
  # ═══════════════════════════════════════════════════════════════════════════
  local -a SELECTED_INDICES=()
  local install_nvidia_drivers="yes"

  # Picks a driver stack, not a card: two NVIDIA cards install the same
  # libraries, and the host already asked which to pass through. lspci sees every
  # GPU regardless, so this asked twice. A mixed set is still a real choice.
  local unique_types
  unique_types="$(printf '%s\n' "${GPU_TYPES[@]}" | sort -u | wc -l)"

  if [[ $gpu_count -eq 1 ]]; then
    # Single GPU - auto-select
    SELECTED_INDICES=(0)
    msg_ok "Detected GPU: ${GPU_NAMES[0]} (${GPU_TYPES[0]})"
  elif [[ "$unique_types" -eq 1 ]]; then
    for i in "${!GPU_LIST[@]}"; do
      SELECTED_INDICES+=("$i")
    done
    msg_ok "Detected ${gpu_count}x ${GPU_TYPES[0]} GPU - configuring all, same driver stack"
  else
    # Multiple GPUs - show selection menu
    echo ""
    msg_custom "⚠" "${YW}" "Multiple GPUs detected:"
    echo ""
    for i in "${!GPU_LIST[@]}"; do
      local type_display="${GPU_TYPES[$i]}"
      case "${GPU_TYPES[$i]}" in
      INTEL_ARC) type_display="Intel Arc" ;;
      INTEL_GEN9+) type_display="Intel Gen9+" ;;
      INTEL_LEGACY) type_display="Intel Legacy" ;;
      INTEL) type_display="Intel" ;;
      AMD) type_display="AMD" ;;
      AMD_APU) type_display="AMD APU" ;;
      NVIDIA) type_display="NVIDIA" ;;
      esac
      printf "  %d) [%s] %s\n" "$((i + 1))" "$type_display" "${GPU_NAMES[$i]}"
    done
    printf "  A) Configure ALL GPUs\n"
    echo ""

    # Read with 60 second timeout
    local selection=""
    echo -n "Select GPU(s) to configure (1-${gpu_count}, A=all) [timeout 60s, default=all]: "
    if read -r -t 60 selection </dev/tty; then
      selection="${selection^^}" # uppercase
    else
      echo ""
      msg_custom "ℹ️" "${YW}" "Timeout - configuring all GPUs automatically"
      selection="A"
    fi

    # Parse selection
    if [[ "$selection" == "A" || -z "$selection" ]]; then
      # Select all
      for i in "${!GPU_LIST[@]}"; do
        SELECTED_INDICES+=("$i")
      done
    elif [[ "$selection" =~ ^[0-9,]+$ ]]; then
      # Parse comma-separated numbers
      IFS=',' read -ra nums <<<"$selection"
      for num in "${nums[@]}"; do
        num=$(echo "$num" | tr -d ' ')
        if [[ "$num" =~ ^[0-9]+$ ]] && ((num >= 1 && num <= gpu_count)); then
          SELECTED_INDICES+=("$((num - 1))")
        fi
      done
    else
      # Invalid - default to all
      msg_warn "Invalid selection - configuring all GPUs"
      for i in "${!GPU_LIST[@]}"; do
        SELECTED_INDICES+=("$i")
      done
    fi
  fi

  # Ask whether to install NVIDIA drivers in the container
  local nvidia_selected="no"
  for idx in "${SELECTED_INDICES[@]}"; do
    if [[ "${GPU_TYPES[$idx]}" == "NVIDIA" ]]; then
      nvidia_selected="yes"
      break
    fi
  done

  if [[ "$nvidia_selected" == "yes" ]]; then
    if [[ -n "${INSTALL_NVIDIA_DRIVERS:-}" ]]; then
      install_nvidia_drivers="${INSTALL_NVIDIA_DRIVERS}"
    else
      echo ""
      msg_custom "🎮" "${GN}" "NVIDIA GPU passthrough detected"
      local nvidia_reply=""
      read -r -t 60 -p "${TAB3}⚙️ Install NVIDIA driver libraries in the container? [Y/n] (auto-yes in 60s): " nvidia_reply </dev/tty || nvidia_reply=""
      case "${nvidia_reply,,}" in
      n | no) install_nvidia_drivers="no" ;;
      *) install_nvidia_drivers="yes" ;;
      esac
    fi
  fi

  # ═══════════════════════════════════════════════════════════════════════════
  # OS Detection
  # ═══════════════════════════════════════════════════════════════════════════
  local os_id os_codename os_version
  os_id=$(get_os_info id)
  os_codename=$(get_os_info codename)
  os_version=$(get_os_info version)
  [[ -z "$os_id" ]] && os_id="debian"

  local in_ct="${CTTYPE:-0}"

  # ═══════════════════════════════════════════════════════════════════════════
  # Process Selected GPUs
  # ═══════════════════════════════════════════════════════════════════════════
  for idx in "${SELECTED_INDICES[@]}"; do
    local gpu_type="${GPU_TYPES[$idx]}"
    local gpu_name="${GPU_NAMES[$idx]}"

    msg_info "Configuring: ${gpu_name}"

    case "$gpu_type" in
    # ─────────────────────────────────────────────────────────────────────────
    # Intel Arc GPUs (DG1, DG2, Arc A-series)
    # ─────────────────────────────────────────────────────────────────────────
    INTEL_ARC)
      _setup_intel_arc "$os_id" "$os_codename"
      ;;

    # ─────────────────────────────────────────────────────────────────────────
    # Intel Gen 9+ (Skylake 2015+: UHD, Iris, HD 6xx+)
    # ─────────────────────────────────────────────────────────────────────────
    INTEL_GEN9+ | INTEL)
      _setup_intel_modern "$os_id" "$os_codename"
      ;;

    # ─────────────────────────────────────────────────────────────────────────
    # Intel Legacy (Gen 6-8: HD 2000-5999, Sandy Bridge to Broadwell)
    # ─────────────────────────────────────────────────────────────────────────
    INTEL_LEGACY)
      _setup_intel_legacy "$os_id" "$os_codename"
      ;;

    # ─────────────────────────────────────────────────────────────────────────
    # AMD Discrete GPUs
    # ─────────────────────────────────────────────────────────────────────────
    AMD)
      _setup_amd_gpu "$os_id" "$os_codename"
      ;;

    # ─────────────────────────────────────────────────────────────────────────
    # AMD APU (Integrated Graphics)
    # ─────────────────────────────────────────────────────────────────────────
    AMD_APU)
      _setup_amd_apu "$os_id" "$os_codename"
      ;;

    # ─────────────────────────────────────────────────────────────────────────
    # NVIDIA GPUs
    # ─────────────────────────────────────────────────────────────────────────
    NVIDIA)
      if [[ "$install_nvidia_drivers" == "yes" ]]; then
        _setup_nvidia_gpu "$os_id" "$os_codename" "$os_version"
      else
        msg_warn "Skipping NVIDIA driver installation (user opted to install manually)"
      fi
      ;;
    esac
  done

  # ═══════════════════════════════════════════════════════════════════════════
  # Device Permissions
  # ═══════════════════════════════════════════════════════════════════════════
  _setup_gpu_permissions "$in_ct" "$service_user"

  local _gpu_type
  for idx in "${SELECTED_INDICES[@]}"; do
    _gpu_type="${GPU_TYPES[$idx]}"
    if [[ "$_gpu_type" == "NVIDIA" && $_cs_nvidia_usable -eq 1 ]]; then
      HWACCEL_VENDOR="nvidia"
      break
    fi
    if [[ "$HWACCEL_VENDOR" == "none" ]]; then
      case "$_gpu_type" in
      AMD) HWACCEL_VENDOR="amd" ;;
      # Not amd: consumers read that as ROCm, and an APU gets Mesa only.
      AMD_APU) HWACCEL_VENDOR="amd_apu" ;;
      INTEL*) HWACCEL_VENDOR="intel" ;;
      esac
    fi
  done
  export HWACCEL_VENDOR

  cache_installed_version "hwaccel" "1.0"
  if [[ "$HWACCEL_VENDOR" == "none" ]]; then
    # Close the step like msg_ok, or a later failure is reported under it.
    # shellcheck disable=SC2034
    CS_STEP_DONE="Setup Hardware Acceleration" CS_STEP="" CS_ERROR_MSG=""
    _cs_drop_silent_record
    msg_warn "Hardware acceleration is not available - see the messages above"
  else
    msg_ok "Setup Hardware Acceleration"
  fi
}

# ══════════════════════════════════════════════════════════════════════════════
# Resolve the IGC tag that the latest compute-runtime was built against.
# Must be called AFTER a fetch_and_deploy_gh_release for intel/compute-runtime
# so that TOOLS_GH_REL_JSON contains the compute-runtime release metadata.
# Sets the variable named by $1 (default: igc_tag) to the discovered tag.
# ══════════════════════════════════════════════════════════════════════════════
# compute-runtime pins the exact intel-igc-opencl-2 it was built against, so the
# tag has to come from its release notes. Empty means unresolved: installing the
# newest IGC instead gives a combination dpkg refuses.
_resolve_igc_tag() {
  local -n _out_ref="${1:-igc_tag}"
  _out_ref=""
  [[ -n "${TOOLS_GH_REL_JSON:-}" && -f "$TOOLS_GH_REL_JSON" ]] || return 0

  # The file holds whatever the API returned: one release for a pinned version,
  # a list for 'latest'.
  local _body
  _body=$(jq -r 'if type == "array" then .[0] else . end | .body // empty' "$TOOLS_GH_REL_JSON" 2>/dev/null) || return 0
  _out_ref=$(grep -oP 'intel-graphics-compiler/releases/tag/\K[^\s\)]+' <<<"$_body" | head -1)
}

# ══════════════════════════════════════════════════════════════════════════════
# Intel Arc GPU Setup
# ══════════════════════════════════════════════════════════════════════════════
_setup_intel_arc() {
  local os_id="$1" os_codename="$2"

  msg_info "Installing Intel Arc GPU drivers"

  if [[ "$os_id" == "ubuntu" ]]; then
    # Ubuntu 22.04+ has Arc support in HWE kernel
    _cs_apt_install_optional intel-media-va-driver-non-free intel-opencl-icd libmfx-gen1.2 mesa-vulkan-drivers vulkan-tools vainfo intel-gpu-tools

    # The oneVPL runtime dispatcher. Anything that dlopens libvpl instead of
    # going through VA-API -- ErsatzTV's QSV capability probe, for one -- fails
    # at runtime without it. On its own line rather than in the batch above: an
    # apt list is all-or-nothing, so a release that does not carry the package
    # would take every other driver down with it.
    $STD apt -y install libvpl2 2>/dev/null || true

  elif [[ "$os_id" == "debian" ]]; then
    # Add non-free repos
    _add_debian_nonfree "$os_codename"

    # For Trixie/Sid: Fetch latest drivers from GitHub (Debian repo packages may be too old or missing)
    # For Bookworm: Use repo packages (GitHub latest requires libstdc++6 >= 13.1, unavailable on Bookworm)
    if [[ "$os_codename" == "trixie" || "$os_codename" == "sid" ]]; then
      msg_custom "ℹ️" "${YW}" "Fetching Intel compute-runtime from GitHub for Arc support"

      # intel-opencl-icd declares it; dpkg refuses the package without it.
      _cs_apt_install_optional ocl-icd-libopencl1

      # Fetch a compute-runtime package first so TOOLS_GH_REL_JSON is populated,
      # then resolve the matching IGC tag from the release notes.
      # libigdgmm - bundled in compute-runtime releases
      fetch_and_deploy_gh_release "libigdgmm12" "intel/compute-runtime" "binary" "latest" "" "libigdgmm12_*_amd64.deb" || true

      local igc_tag
      _resolve_igc_tag igc_tag
      if [[ -z "$igc_tag" ]]; then
        msg_warn "Could not read the IGC version compute-runtime pins - skipping OpenCL. VA-API still works."
      else

        # Intel Graphics Compiler – pinned to the version compute-runtime expects
        fetch_and_deploy_gh_release "intel-igc-core" "intel/intel-graphics-compiler" "binary" "$igc_tag" "" "intel-igc-core-2_*_amd64.deb" || true
        fetch_and_deploy_gh_release "intel-igc-opencl" "intel/intel-graphics-compiler" "binary" "$igc_tag" "" "intel-igc-opencl-2_*_amd64.deb" || true

        # Compute Runtime (depends on IGC and gmmlib)
        fetch_and_deploy_gh_release "intel-opencl-icd" "intel/compute-runtime" "binary" "latest" "" "intel-opencl-icd_*_amd64.deb" || true
        fetch_and_deploy_gh_release "intel-level-zero-gpu" "intel/compute-runtime" "binary" "latest" "" "libze-intel-gpu1_*_amd64.deb" || true
      fi
    fi

    _cs_apt_install_optional intel-media-va-driver-non-free ocl-icd-libopencl1 libvpl2 libmfx-gen1.2 mesa-vulkan-drivers vulkan-tools vainfo intel-gpu-tools

    # Bookworm has compatible versions of these packages in repos
    [[ "$os_codename" == "bookworm" ]] && _cs_apt_install_optional intel-opencl-icd libigdgmm12
  fi

  msg_ok "Intel Arc GPU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# Intel Modern GPU Setup (Gen 9+)
# ══════════════════════════════════════════════════════════════════════════════
_setup_intel_modern() {
  local os_id="$1" os_codename="$2"

  msg_info "Installing Intel Gen 9+ GPU drivers"

  if [[ "$os_id" == "ubuntu" ]]; then
    _cs_apt_install_optional va-driver-all intel-media-va-driver ocl-icd-libopencl1 mesa-vulkan-drivers vulkan-tools vainfo intel-gpu-tools

    # Try non-free driver for better codec support
    $STD apt -y install intel-media-va-driver-non-free 2>/dev/null || true
    $STD apt -y install intel-opencl-icd 2>/dev/null || true
    $STD apt -y install libmfx-gen1.2 2>/dev/null || true
    # oneVPL dispatcher — see _setup_intel_arc for why this matters
    $STD apt -y install libvpl2 2>/dev/null || true

  elif [[ "$os_id" == "debian" ]]; then
    _add_debian_nonfree "$os_codename"

    # For Trixie/Sid: Fetch from GitHub (Debian packages too old or missing)
    if [[ "$os_codename" == "trixie" || "$os_codename" == "sid" ]]; then
      msg_custom "ℹ️" "${YW}" "Fetching Intel compute-runtime from GitHub"

      # intel-opencl-icd declares it; dpkg refuses the package without it.
      _cs_apt_install_optional ocl-icd-libopencl1

      # Fetch a compute-runtime package first so TOOLS_GH_REL_JSON is populated,
      # then resolve the matching IGC tag from the release notes.
      # libigdgmm first (bundled in compute-runtime releases)
      fetch_and_deploy_gh_release "libigdgmm12" "intel/compute-runtime" "binary" "latest" "" "libigdgmm12_*_amd64.deb" || true

      local igc_tag
      _resolve_igc_tag igc_tag
      if [[ -z "$igc_tag" ]]; then
        msg_warn "Could not read the IGC version compute-runtime pins - skipping OpenCL. VA-API still works."
      else

        # Intel Graphics Compiler – pinned to the version compute-runtime expects
        fetch_and_deploy_gh_release "intel-igc-core" "intel/intel-graphics-compiler" "binary" "$igc_tag" "" "intel-igc-core-2_*_amd64.deb" || true
        fetch_and_deploy_gh_release "intel-igc-opencl" "intel/intel-graphics-compiler" "binary" "$igc_tag" "" "intel-igc-opencl-2_*_amd64.deb" || true

        # Compute Runtime
        fetch_and_deploy_gh_release "intel-opencl-icd" "intel/compute-runtime" "binary" "latest" "" "intel-opencl-icd_*_amd64.deb" || true
      fi
    fi

    _cs_apt_install_optional intel-media-va-driver-non-free ocl-icd-libopencl1 mesa-vulkan-drivers vulkan-tools vainfo libmfx-gen1.2 intel-gpu-tools

    # oneVPL dispatcher — see _setup_intel_arc for why this matters
    $STD apt -y install libvpl2 2>/dev/null || true

    # Bookworm has intel-opencl-icd in repos (compatible version)
    [[ "$os_codename" == "bookworm" ]] && _cs_apt_install_optional intel-opencl-icd libigdgmm12
  fi

  msg_ok "Intel Gen 9+ GPU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# Intel Legacy GPU Setup (Gen 6-8)
# ══════════════════════════════════════════════════════════════════════════════
_setup_intel_legacy() {
  local os_id="$1" os_codename="$2"

  msg_info "Installing Intel Legacy GPU drivers (Gen 6-8)"

  # Legacy GPUs use i965 driver - stable repo packages only
  _cs_apt_install_optional va-driver-all i965-va-driver mesa-va-drivers ocl-icd-libopencl1 vainfo intel-gpu-tools

  # beignet provides OpenCL for older Intel GPUs (Sandy Bridge to Broadwell)
  # Note: beignet-opencl-icd was removed in Debian 12+ and Ubuntu 22.04+
  # Check if package is available before attempting installation
  if apt-cache show beignet-opencl-icd &>/dev/null; then
    $STD apt -y install beignet-opencl-icd 2>/dev/null || msg_warn "beignet-opencl-icd installation failed (optional)"
  else
    msg_warn "beignet-opencl-icd not available - OpenCL support for legacy Intel GPU limited"
    msg_warn "Note: Hardware video encoding/decoding (VA-API) still works without OpenCL"
  fi

  msg_ok "Intel Legacy GPU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# AMD Discrete GPU Setup
# ══════════════════════════════════════════════════════════════════════════════
_setup_amd_gpu() {
  local os_id="$1" os_codename="$2"

  msg_info "Installing AMD GPU drivers"

  # Core Mesa drivers
  _cs_apt_install_optional mesa-va-drivers mesa-vdpau-drivers mesa-opencl-icd ocl-icd-libopencl1 libdrm-amdgpu1 vainfo clinfo

  # Firmware for AMD GPUs
  if [[ "$os_id" == "debian" ]]; then
    _add_debian_nonfree_firmware "$os_codename"
    $STD apt -y install firmware-amd-graphics 2>/dev/null || msg_warn "AMD firmware not available"
  fi
  # Ubuntu includes AMD firmware in linux-firmware by default

  # ROCm compute stack (OpenCL + HIP)
  _setup_rocm "$os_id" "$os_codename"

  msg_ok "AMD GPU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# AMD APU Setup (Integrated Graphics)
# ══════════════════════════════════════════════════════════════════════════════
_setup_amd_apu() {
  local os_id="$1" os_codename="$2"

  msg_info "Installing AMD APU drivers"

  _cs_apt_install_optional mesa-va-drivers mesa-vdpau-drivers mesa-opencl-icd ocl-icd-libopencl1 vainfo

  if [[ "$os_id" == "debian" ]]; then
    _add_debian_nonfree_firmware "$os_codename"
    $STD apt -y install firmware-amd-graphics 2>/dev/null || true
  fi

  msg_ok "AMD APU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# AMD ROCm Compute Setup
# Adds ROCm repository and installs the ROCm compute stack for AMD GPUs/APUs.
# Provides: OpenCL, HIP, rocm-smi, rocminfo
# Supported: Debian 12/13, Ubuntu 22.04/24.04 (amd64 only)
# ══════════════════════════════════════════════════════════════════════════════
_setup_rocm() {
  local os_id="$1" os_codename="$2"

  # Only amd64 is supported
  if [[ "$(dpkg --print-architecture 2>/dev/null)" != "amd64" ]]; then
    msg_warn "ROCm is only available for amd64 — skipping"
    return 0
  fi

  local ROCM_VERSION="7.2"
  local ROCM_REPO_CODENAME

  # Map OS codename to ROCm repository codename (Ubuntu-based repos)
  case "${os_id}-${os_codename}" in
  debian-bookworm) ROCM_REPO_CODENAME="jammy" ;;
  debian-trixie | debian-sid) ROCM_REPO_CODENAME="noble" ;;
  ubuntu-jammy) ROCM_REPO_CODENAME="jammy" ;;
  ubuntu-noble) ROCM_REPO_CODENAME="noble" ;;
  *)
    msg_warn "ROCm not supported on ${os_id} ${os_codename} — skipping"
    return 0
    ;;
  esac

  msg_info "Installing ROCm ${ROCM_VERSION} compute stack"

  # ROCm main repository (userspace compute libs)
  setup_deb822_repo \
    "rocm" \
    "https://repo.radeon.com/rocm/rocm.gpg.key" \
    "https://repo.radeon.com/rocm/apt/${ROCM_VERSION}" \
    "${ROCM_REPO_CODENAME}" \
    "main" \
    "amd64" || {
    msg_warn "Failed to add ROCm repository — skipping ROCm"
    return 0
  }

  # Note: The amdgpu/latest/ubuntu repo (kernel driver packages) is intentionally
  # omitted — kernel drivers are managed by the Proxmox host, not the LXC container.
  # Only the ROCm userspace compute stack is needed inside the container.

  # Pin ROCm packages to prefer radeon repo
  cat <<EOF >/etc/apt/preferences.d/rocm-pin-600
Package: *
Pin: release o=repo.radeon.com
Pin-Priority: 600
EOF

  # apt update with retry — repo.radeon.com CDN can be mid-sync (transient size mismatches).
  # Run with ERR trap disabled so a transient failure does not abort the entire install.
  local _apt_ok=0
  for _attempt in 1 2 3; do
    if (
      set +e
      apt-get update -qq 2>&1
      exit $?
    ) 2>/dev/null; then
      _apt_ok=1
      break
    fi
    msg_warn "apt update failed (attempt ${_attempt}/3) — AMD repo may be temporarily unavailable, retrying in 30s…"
    sleep 30
  done
  if [[ $_apt_ok -eq 0 ]]; then
    msg_warn "apt update still failing after 3 attempts — skipping ROCm install"
    return 0
  fi

  # Install only runtime packages — full 'rocm' meta-package includes 15GB+ dev tools
  $STD apt install -y rocm-opencl-runtime rocm-hip-runtime rocm-smi-lib 2>/dev/null || {
    msg_warn "ROCm runtime install failed — trying minimal set"
    _cs_apt_install_optional rocm-opencl-runtime rocm-smi-lib
  }

  # Group membership for GPU access
  usermod -aG render,video root 2>/dev/null || true

  # Environment (PATH + LD_LIBRARY_PATH)
  if [[ -d /opt/rocm ]]; then
    cat <<'ENVEOF' >/etc/profile.d/rocm.sh
export PATH="$PATH:/opt/rocm/bin"
export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:+$LD_LIBRARY_PATH:}/opt/rocm/lib"
ENVEOF
    chmod +x /etc/profile.d/rocm.sh
    # Also make available for current session / systemd services
    echo "/opt/rocm/lib" >/etc/ld.so.conf.d/rocm.conf
    ldconfig 2>/dev/null || true
  fi

  if [[ -x /opt/rocm/bin/rocminfo ]]; then
    msg_ok "ROCm ${ROCM_VERSION} installed"
  else
    msg_warn "ROCm installed but rocminfo not found — GPU may not be available in container"
  fi
}

# ══════════════════════════════════════════════════════════════════════════════
# NVIDIA GPU Setup
# ══════════════════════════════════════════════════════════════════════════════

# The host's exact release first; the series fallback may end up newer than
# the kernel module, which the caller warns about.
_cuda_pick_version() {
  local host="$1" madison exact
  madison=$(apt-cache madison libcuda1 2>/dev/null | awk '{print $3}') || return 0
  exact=$(printf '%s\n' "$madison" | grep -E "^${host//./\\.}(-|$)" | head -1 || true)
  if [[ -n "$exact" ]]; then
    printf '%s\n' "$exact"
    return 0
  fi
  printf '%s\n' "$madison" | grep -E "^${host%%.*}\." | head -1 || true
}

# libcuda1 needs its libnvidia-* dependencies at exactly its own version.
# The priority has to stay above the origin pin's.
_nvidia_version_pin() {
  cat <<NVIDIA_VERSION_PIN
Package: libcuda* libnvcuvid* libnvidia-*
Pin: version $1
Pin-Priority: 1002
NVIDIA_VERSION_PIN
}

# Version stanza first: apt uses the first stanza that matches a version.
# The origin pin stays scoped, since above 1000 apt may downgrade what it matches.
_nvidia_apt_pins() {
  if [[ -n "${1:-}" ]]; then
    _nvidia_version_pin "$1"
    echo
  fi
  cat <<'NVIDIA_PIN'
Package: libcuda* libnvcuvid* libnvidia-* nvidia-smi nvidia-utils-* cuda-*
Pin: origin developer.download.nvidia.com
Pin-Priority: 1001
NVIDIA_PIN
}

_setup_nvidia_gpu() {
  local os_id="$1" os_codename="$2" os_version="$3"

  msg_info "Installing NVIDIA GPU drivers"

  # Prevent interactive dialogs (e.g., "Mismatching nvidia kernel module" whiptail)
  export DEBIAN_FRONTEND=noninteractive
  export NEEDRESTART_MODE=a

  # Detect host driver version (passed through via /proc)
  # Format varies by driver type:
  #   Proprietary: "NVRM version: NVIDIA UNIX x86_64 Kernel Module  550.54.14  Thu..."
  #   Open:        "NVRM version: NVIDIA UNIX Open Kernel Module for x86_64  590.48.01  Release..."
  # Use regex to extract version number (###.##.## or ###.## pattern)
  local nvidia_host_version=""
  if [[ -f /proc/driver/nvidia/version ]]; then
    nvidia_host_version=$(grep -oP '\d{3,}\.\d+(\.\d+)?' /proc/driver/nvidia/version 2>/dev/null | head -1 || true)
  fi

  if [[ -z "$nvidia_host_version" ]]; then
    msg_warn "NVIDIA host driver version not found in /proc/driver/nvidia/version"
    msg_warn "Ensure NVIDIA drivers are installed on host and GPU passthrough is enabled"
    _cs_apt_install_optional va-driver-all vainfo
    return 0
  fi

  msg_custom "ℹ️" "${YW}" "Host NVIDIA driver version: ${nvidia_host_version}"

  if [[ "$os_id" == "debian" ]]; then
    # Enable non-free components
    if [[ -f /etc/apt/sources.list.d/debian.sources ]]; then
      if ! grep -q "non-free" /etc/apt/sources.list.d/debian.sources 2>/dev/null; then
        sed -i -E 's/Components: (.*)$/Components: \1 contrib non-free non-free-firmware/g' /etc/apt/sources.list.d/debian.sources 2>/dev/null || true
      fi
    fi
    $STD apt-get -y update 2>/dev/null || msg_warn "apt update failed - continuing anyway"

    # Which CUDA repository can actually carry this host's driver?
    #
    # Choosing by codename alone is wrong, and a user's report proved it. NVIDIA
    # does publish a debian13 repo, but it starts at 590: a host on 580 finds
    # nothing there at all, while debian12 carries 545..610 including 580. The
    # old code sent trixie to Debian's own repos for the same reason and found
    # nothing either, so a Debian 13 container ended up with no NVIDIA userspace
    # whatsoever -- dpkg -l | grep nvidia came back empty on the reporter's host.
    #
    # "Newest repo wins" is not a safe default in the other direction either:
    # libcuda is the driver library and has to agree with the loaded kernel
    # module, so a userspace newer than the host cannot work.
    #
    # The codename therefore only sets the ORDER of candidates. Which one is
    # used is decided by whether it actually carries the host's driver series.
    # Packages built for bookworm load on trixie -- they link against glibc,
    # which is backward compatible.
    local cuda_repos=()
    case "$os_codename" in
    bullseye) cuda_repos=(debian11 debian12) ;;
    bookworm) cuda_repos=(debian12) ;;
    trixie | sid) cuda_repos=(debian13 debian12) ;;
    *) cuda_repos=(debian12) ;;
    esac

    # One key signs every CUDA repo, so the keyring is fetched once regardless
    # of which candidate wins.
    if [[ ! -f /usr/share/keyrings/cuda-archive-keyring.gpg ]]; then
      msg_info "Adding NVIDIA CUDA repository keyring"
      local cuda_keyring
      cuda_keyring="$(mktemp)"
      if curl -fsSL -o "$cuda_keyring" "https://developer.download.nvidia.com/compute/cuda/repos/${cuda_repos[0]}/x86_64/cuda-keyring_1.1-1_all.deb" 2>/dev/null; then
        $STD dpkg -i "$cuda_keyring" 2>/dev/null || true
        msg_ok "Added NVIDIA CUDA repository keyring"
      else
        msg_warn "Failed to download NVIDIA CUDA keyring"
      fi
      rm -f "$cuda_keyring"
    fi

    local nvidia_pin_file="/etc/apt/preferences.d/nvidia-cuda-pin"
    [[ -f "$nvidia_pin_file" ]] || _nvidia_apt_pins >"$nvidia_pin_file"

    # Extract major version for flexible matching (580.126.09 -> 580)
    local nvidia_major_version="${nvidia_host_version%%.*}"
    local cuda_list="/etc/apt/sources.list.d/cuda-community-scripts.list"
    local chosen_repo="" available_version="" cuda_repo=""

    for cuda_repo in "${cuda_repos[@]}"; do
      # The keyring package drops its own list file for whichever repo it came
      # from. Left in place it would keep a rejected candidate active alongside
      # the one being tried, so only this file decides.
      rm -f /etc/apt/sources.list.d/cuda-debian*.list
      echo "deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] https://developer.download.nvidia.com/compute/cuda/repos/${cuda_repo}/x86_64/ /" >"$cuda_list"
      $STD apt-get -y update 2>/dev/null || true

      available_version="$(_cuda_pick_version "$nvidia_host_version")"
      if [[ -n "$available_version" ]]; then
        chosen_repo="$cuda_repo"
        break
      fi
      msg_warn "CUDA repository ${cuda_repo} carries no ${nvidia_major_version}.x driver"
    done

    if [[ -z "$chosen_repo" ]]; then
      # Last resort: the distribution's own non-free packages.
      rm -f "$cuda_list"
      $STD apt-get -y update 2>/dev/null || true
      available_version="$(_cuda_pick_version "$nvidia_host_version")"
      [[ -n "$available_version" ]] && chosen_repo="Debian non-free"
    fi

    if [[ -n "$chosen_repo" ]]; then
      _nvidia_apt_pins "$available_version" >"$nvidia_pin_file"
      msg_info "Installing NVIDIA userspace ${available_version} from ${chosen_repo}"
      local nvidia_pkgs="libcuda1=${available_version} libnvcuvid1=${available_version} libnvidia-encode1=${available_version} libnvidia-ml1=${available_version}"
      # -y refuses the downgrade a host rollback needs, pin or not.
      local -a apt_install=(apt-get -y --allow-downgrades -o Dpkg::Options::="--force-confold" install --no-install-recommends)
      if $STD "${apt_install[@]}" $nvidia_pkgs 2>/dev/null; then
        # Say which kind of match this was. NVIDIA wants the full version to
        # agree; the fallback above only guarantees the series, so 580.178.04
        # on the host and 580.126.20 here would both land here. State both and
        # let the reader judge rather than calling it "version-matched".
        if [[ "$available_version" == "${nvidia_host_version}"* ]]; then
          msg_ok "NVIDIA userspace ${available_version} matches the host driver"
        else
          msg_warn "NVIDIA userspace ${available_version} differs from the host driver ${nvidia_host_version}"
          msg_warn "Same driver series, different release. CUDA workloads may still fail on this."
        fi
      else
        # A partly pinned retry cannot resolve either; report apt's reason.
        msg_error "NVIDIA userspace ${available_version} would not install from ${chosen_repo}"
        local apt_reason="" reason_line
        apt_reason="$("${apt_install[@]}" -s $nvidia_pkgs 2>&1 | grep -E 'Depends:|^E: ' | sed 's/^ *//' | head -3 || true)"
        while IFS= read -r reason_line; do
          if [[ -n "$reason_line" ]]; then msg_warn "apt: ${reason_line}"; fi
        done <<<"$apt_reason"
      fi
    else
      # Deliberately install nothing. The previous code fell back to whatever
      # version the repo happened to hold, which for a 580 host on trixie would
      # mean a 590+ userspace -- newer than the kernel module, so every CUDA
      # call fails at initialisation. A clear message beats libraries that
      # cannot work.
      msg_warn "No NVIDIA ${nvidia_major_version}.x userspace libraries exist for ${os_id} ${os_codename}"
      msg_warn "Searched: ${cuda_repos[*]} and the distribution's own repositories"
      msg_warn "Host driver is ${nvidia_host_version}. Either run a driver series the container OS carries, or pick a container OS that carries this one."
    fi

    $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends nvidia-smi 2>/dev/null || true

  elif [[ "$os_id" == "ubuntu" ]]; then
    # Ubuntu versioning
    local ubuntu_cuda_repo=""
    case "$os_version" in
    22.04) ubuntu_cuda_repo="ubuntu2204" ;;
    24.04) ubuntu_cuda_repo="ubuntu2404" ;;
    *) ubuntu_cuda_repo="ubuntu2204" ;; # Fallback
    esac

    # Add NVIDIA CUDA repository for Ubuntu
    if [[ ! -f /usr/share/keyrings/cuda-archive-keyring.gpg ]]; then
      msg_info "Adding NVIDIA CUDA repository (${ubuntu_cuda_repo})"
      local cuda_keyring
      cuda_keyring="$(mktemp)"
      if curl -fsSL -o "$cuda_keyring" "https://developer.download.nvidia.com/compute/cuda/repos/${ubuntu_cuda_repo}/x86_64/cuda-keyring_1.1-1_all.deb" 2>/dev/null; then
        $STD dpkg -i "$cuda_keyring" 2>/dev/null || true
        msg_ok "Added NVIDIA CUDA repository (${ubuntu_cuda_repo})"
      else
        msg_warn "Failed to download NVIDIA CUDA keyring"
      fi
      rm -f "$cuda_keyring"
    fi

    $STD apt-get -y update 2>/dev/null || msg_warn "apt update failed - continuing anyway"

    # Extract major version for flexible matching
    local nvidia_major_version="${nvidia_host_version%%.*}"

    # Check what versions are available
    local available_version=""
    available_version=$(apt-cache madison libcuda1 2>/dev/null | awk '{print $3}' | grep -E "^${nvidia_major_version}\." | head -1 || true)

    if [[ -n "$available_version" ]]; then
      msg_info "Installing NVIDIA libraries (version ${available_version})"
      local nvidia_pkgs="libcuda1=${available_version} libnvcuvid1=${available_version} libnvidia-encode1=${available_version} libnvidia-ml1=${available_version}"
      if $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends $nvidia_pkgs 2>/dev/null; then
        # Same as above: only the major was compared.
        if [[ "$available_version" == "${nvidia_host_version}"* ]]; then
          msg_ok "Installed NVIDIA libraries ${available_version} (host ${nvidia_host_version})"
        else
          msg_warn "NVIDIA libraries ${available_version} do not match the host driver ${nvidia_host_version}"
          msg_warn "CUDA workloads usually fail on a mismatch. Install the matching driver on the host, or pick a container OS whose repos carry it."
        fi
      else
        # Fallback to Ubuntu repo packages with versioned nvidia-utils
        msg_warn "CUDA repo install failed - trying Ubuntu native packages (nvidia-utils-${nvidia_major_version})"
        if $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends \
          libnvidia-decode-${nvidia_major_version} libnvidia-encode-${nvidia_major_version} nvidia-utils-${nvidia_major_version} 2>/dev/null; then
          msg_ok "Installed Ubuntu NVIDIA packages (${nvidia_major_version})"
        else
          msg_warn "NVIDIA driver installation failed - please install manually: apt install nvidia-utils-${nvidia_major_version}"
        fi
      fi
    else
      msg_warn "No NVIDIA packages for version ${nvidia_major_version}.x in CUDA repo"
      # Fallback to Ubuntu repo packages with versioned nvidia-utils
      msg_info "Trying Ubuntu native packages (nvidia-utils-${nvidia_major_version})"
      if $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends \
        libnvidia-decode-${nvidia_major_version} libnvidia-encode-${nvidia_major_version} nvidia-utils-${nvidia_major_version} 2>/dev/null; then
        msg_ok "Installed Ubuntu NVIDIA packages (${nvidia_major_version})"
      else
        msg_warn "NVIDIA driver installation failed - please install manually: apt install nvidia-utils-${nvidia_major_version}"
      fi
    fi
  fi

  # VA-API for hybrid setups (Intel + NVIDIA)
  _cs_apt_install_optional va-driver-all vainfo

  # Fix GLX alternatives: nvidia-alternative diverts mesa libs but in LXC
  # containers the nvidia GLX libs are typically missing, leaving libGL.so.1
  # pointing nowhere. Fall back to mesa if nvidia GLX dir is empty/missing.
  if command -v update-glx &>/dev/null; then
    local nvidia_glx_dir="/usr/lib/nvidia"
    if [[ ! -f "${nvidia_glx_dir}/libGL.so.1" ]] && [[ -d /usr/lib/mesa-diverted ]]; then
      msg_custom "ℹ️" "${YW}" "NVIDIA GLX libs missing in container - falling back to mesa"
      $STD update-glx --set glx /usr/lib/mesa-diverted 2>/dev/null || true
      ldconfig 2>/dev/null || true
    fi
  fi

  # Report what happened, not that the function reached its end -- this printed
  # "NVIDIA GPU configured" even when every path above had found nothing.
  # ldconfig rather than dpkg: the library can arrive by bind mount too, and what
  # matters is whether a CUDA program can link it. Not a non-zero return, since
  # acceleration is optional and install scripts run under an ERR trap.
  # Libraries alone are not enough, and this is the half that used to be missed:
  # lspci sees every GPU on the host whether or not it was passed through, so
  # this function will happily install the userspace stack for a card the
  # container cannot reach. That reported success while torch.cuda.is_available()
  # stayed False, and the real fault -- no device nodes -- went unmentioned.
  local -a nvidia_nodes=()
  local _nvdev
  for _nvdev in /dev/nvidia*; do
    [[ -c "$_nvdev" ]] && nvidia_nodes+=("$_nvdev")
  done

  local have_libcuda=0
  ldconfig -p 2>/dev/null | grep -q 'libcuda\.so\.1' && have_libcuda=1

  local libcuda_version=""
  libcuda_version="$(dpkg-query -W -f='${Version}' libcuda1 2>/dev/null || true)"

  if [[ ${#nvidia_nodes[@]} -eq 0 ]]; then
    msg_error "NVIDIA GPU not usable: no /dev/nvidia* device nodes in this container"
    msg_warn "The userspace driver is in place, but the devices were never passed through."
    msg_warn "On the host, check the container config lists /dev/nvidia0, /dev/nvidiactl and /dev/nvidia-uvm."
  elif [[ $have_libcuda -eq 0 ]]; then
    msg_error "NVIDIA GPU not usable: libcuda.so.1 is not available in this container"
    msg_warn "The host driver is ${nvidia_host_version}. CUDA workloads will fail until a matching userspace library is present."
  elif [[ -n "$libcuda_version" && "${libcuda_version%%.*}" != "${nvidia_host_version%%.*}" ]]; then
    msg_error "NVIDIA GPU not usable: libcuda ${libcuda_version} is not the host driver's ${nvidia_host_version%%.*} series"
    msg_warn "The userspace library must match the host module. Install the ${nvidia_host_version%%.*} series on the host, or run NVIDIA's ${nvidia_host_version} .run installer with --no-kernel-module in the container."
  else
    _cs_nvidia_usable=1
    msg_ok "NVIDIA GPU configured (host driver ${nvidia_host_version}, ${#nvidia_nodes[@]} device nodes)"
  fi

  # nvidia-smi separately: in NVIDIA's CUDA repos the package is an empty
  # transitional stub, so dpkg reports it installed while the binary is absent.
  # CUDA does not need it -- libnvidia-ml1 carries NVML -- so this reports rather
  # than pulling the package that ships it, which would drag the full driver
  # stack, DKMS included, into a container sharing the host kernel.
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    msg_warn "nvidia-smi is not available: the CUDA repository ships it as an empty transitional package"
    msg_warn "CUDA workloads are unaffected. For the tool itself, use the matching .run installer with --no-kernel-module."
  fi
}

# ══════════════════════════════════════════════════════════════════════════════
# Helper: Add Debian non-free repositories
# ══════════════════════════════════════════════════════════════════════════════
_add_debian_nonfree() {
  local os_codename="$1"

  [[ -f /etc/apt/sources.list.d/non-free.sources ]] && return 0

  case "$os_codename" in
  bullseye)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: bullseye bullseye-updates
Components: non-free
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  bookworm)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: bookworm bookworm-updates
Components: non-free non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  trixie | sid)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: trixie trixie-updates
Components: non-free non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg

Types: deb
URIs: http://deb.debian.org/debian-security
Suites: trixie-security
Components: non-free non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  esac
  $STD apt -y update
}

# ══════════════════════════════════════════════════════════════════════════════
# Helper: Add Debian non-free-firmware repository
# ══════════════════════════════════════════════════════════════════════════════
_add_debian_nonfree_firmware() {
  local os_codename="$1"

  [[ -f /etc/apt/sources.list.d/non-free-firmware.sources ]] && return 0

  case "$os_codename" in
  bullseye)
    # Debian 11 uses 'non-free' component (no separate non-free-firmware)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free-firmware.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: bullseye bullseye-updates
Components: non-free
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg

Types: deb
URIs: http://deb.debian.org/debian-security
Suites: bullseye-security
Components: non-free
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  bookworm)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free-firmware.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: bookworm bookworm-updates
Components: non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg

Types: deb
URIs: http://deb.debian.org/debian-security
Suites: bookworm-security
Components: non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  trixie | sid)
    cat <<'EOF' >/etc/apt/sources.list.d/non-free-firmware.sources
Types: deb
URIs: http://deb.debian.org/debian
Suites: trixie trixie-updates
Components: non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg

Types: deb
URIs: http://deb.debian.org/debian-security
Suites: trixie-security
Components: non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ;;
  esac
  $STD apt -y update
}

# ══════════════════════════════════════════════════════════════════════════════
# Helper: Setup GPU device permissions
# ══════════════════════════════════════════════════════════════════════════════
_setup_gpu_permissions() {
  local in_ct="$1"
  local service_user="${2:-}"

  # /dev/dri permissions (Intel/AMD)
  if [[ "$in_ct" == "0" && -d /dev/dri ]]; then
    if ls /dev/dri/card* /dev/dri/renderD* &>/dev/null; then
      chgrp video /dev/dri 2>/dev/null || true
      chmod 755 /dev/dri 2>/dev/null || true
      chmod 660 /dev/dri/* 2>/dev/null || true
      $STD adduser "$(id -u -n)" video 2>/dev/null || true
      $STD adduser "$(id -u -n)" render 2>/dev/null || true

      # Sync GID with host
      local host_video_gid host_render_gid
      host_video_gid=$(getent group video | cut -d: -f3)
      host_render_gid=$(getent group render | cut -d: -f3)
      if [[ -n "$host_video_gid" ]]; then
        sed -i "s/^video:x:[0-9]*:/video:x:$host_video_gid:/" /etc/group 2>/dev/null || true
      fi
      if [[ -n "$host_render_gid" ]]; then
        sed -i "s/^render:x:[0-9]*:/render:x:$host_render_gid:/" /etc/group 2>/dev/null || true
      fi

      # Verify VA-API
      if command -v vainfo &>/dev/null; then
        if vainfo &>/dev/null; then
          msg_custom "✔️" "${GN}" "VA-API verified and working"
        else
          msg_warn "vainfo test failed - check GPU passthrough"
        fi
      fi
    fi
  fi

  # /dev/nvidia* permissions (NVIDIA)
  if ls /dev/nvidia* &>/dev/null 2>&1; then
    msg_info "Configuring NVIDIA device permissions"
    for nvidia_dev in /dev/nvidia*; do
      [[ -e "$nvidia_dev" ]] && {
        chgrp video "$nvidia_dev" 2>/dev/null || true
        chmod 660 "$nvidia_dev" 2>/dev/null || true
      }
    done
    if [[ -d /dev/nvidia-caps ]]; then
      chmod 755 /dev/nvidia-caps 2>/dev/null || true
      for caps_dev in /dev/nvidia-caps/*; do
        [[ -e "$caps_dev" ]] && {
          chgrp video "$caps_dev" 2>/dev/null || true
          chmod 660 "$caps_dev" 2>/dev/null || true
        }
      done
    fi
    msg_ok "Configured NVIDIA device permissions"

    # Verify nvidia-smi
    if command -v nvidia-smi &>/dev/null; then
      if nvidia-smi &>/dev/null; then
        msg_custom "✔️" "${GN}" "nvidia-smi verified and working"
      else
        msg_warn "nvidia-smi test failed - check driver version match"
      fi
    fi
  fi

  # /dev/kfd permissions (AMD ROCm)
  if [[ -e /dev/kfd ]]; then
    chgrp render /dev/kfd 2>/dev/null || true
    chmod 660 /dev/kfd 2>/dev/null || true
    msg_custom "✔️" "${GN}" "AMD ROCm compute device configured"
  fi

  # Add service user to render and video groups for GPU hardware acceleration
  if [[ -n "$service_user" ]]; then
    usermod -aG render "$service_user" 2>/dev/null || true
    usermod -aG video "$service_user" 2>/dev/null || true
  fi
}