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

# ------------------------------------------------------------------------------
# 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:-}"

  # 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
      gpu_type="NVIDIA"
    fi

    if [[ -n "$gpu_type" ]]; then
      GPU_LIST+=("$pci_addr")
      GPU_TYPES+=("$gpu_type")
      GPU_NAMES+=("$gpu_name")
      ((gpu_count++)) || true
    fi
  done < <(lspci 2>/dev/null | grep -Ei 'vga|3d|display')

  # 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"

  if [[ $gpu_count -eq 1 ]]; then
    # Single GPU - auto-select
    SELECTED_INDICES=(0)
    msg_ok "Detected GPU: ${GPU_NAMES[0]} (${GPU_TYPES[0]})"
  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_info "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"

  cache_installed_version "hwaccel" "1.0"
  msg_ok "Setup Hardware Acceleration"
}

# ══════════════════════════════════════════════════════════════════════════════
# 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.
# ══════════════════════════════════════════════════════════════════════════════
_resolve_igc_tag() {
  local -n _out_ref="${1:-igc_tag}"
  _out_ref="latest"
  if [[ -n "${TOOLS_GH_REL_JSON:-}" && -f "$TOOLS_GH_REL_JSON" ]]; then
    local _body _parsed
    _body=$(jq -r '.body // empty' "$TOOLS_GH_REL_JSON" 2>/dev/null) || return 0
    _parsed=$(grep -oP 'intel-graphics-compiler/releases/tag/\K[^\s\)]+' <<<"$_body" | head -1)
    [[ -n "$_parsed" ]] && _out_ref="$_parsed"
  fi
}

# ══════════════════════════════════════════════════════════════════════════════
# 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
    $STD apt -y install \
      intel-media-va-driver-non-free \
      intel-opencl-icd \
      libmfx-gen1.2 \
      mesa-vulkan-drivers \
      vulkan-tools \
      vainfo \
      intel-gpu-tools 2>/dev/null || msg_warn "Some Intel Arc packages failed"

  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_info "Fetching Intel compute-runtime from GitHub for Arc support"

      # 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

      # 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

    $STD apt -y install \
      intel-media-va-driver-non-free \
      ocl-icd-libopencl1 \
      libvpl2 \
      libmfx-gen1.2 \
      mesa-vulkan-drivers \
      vulkan-tools \
      vainfo \
      intel-gpu-tools 2>/dev/null || msg_warn "Some Intel Arc packages failed"

    # Bookworm has compatible versions of these packages in repos
    [[ "$os_codename" == "bookworm" ]] && $STD apt -y install intel-opencl-icd libigdgmm12 2>/dev/null || true
  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
    $STD apt -y install \
      va-driver-all \
      intel-media-va-driver \
      ocl-icd-libopencl1 \
      mesa-vulkan-drivers \
      vulkan-tools \
      vainfo \
      intel-gpu-tools 2>/dev/null || msg_warn "Some Intel packages failed"

    # 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

  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_info "Fetching Intel compute-runtime from GitHub"

      # 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

      # 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

    $STD apt -y install \
      intel-media-va-driver-non-free \
      ocl-icd-libopencl1 \
      mesa-vulkan-drivers \
      vulkan-tools \
      vainfo \
      libmfx-gen1.2 \
      intel-gpu-tools 2>/dev/null || msg_warn "Some Intel packages failed"

    # Bookworm has intel-opencl-icd in repos (compatible version)
    [[ "$os_codename" == "bookworm" ]] && $STD apt -y install intel-opencl-icd libigdgmm12 2>/dev/null || true
  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
  $STD apt -y install \
    va-driver-all \
    i965-va-driver \
    mesa-va-drivers \
    ocl-icd-libopencl1 \
    vainfo \
    intel-gpu-tools 2>/dev/null || msg_warn "Some Intel legacy packages failed"

  # 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
  $STD apt -y install \
    mesa-va-drivers \
    mesa-vdpau-drivers \
    mesa-opencl-icd \
    ocl-icd-libopencl1 \
    libdrm-amdgpu1 \
    vainfo \
    clinfo 2>/dev/null || msg_warn "Some AMD packages failed"

  # 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"

  $STD apt -y install \
    mesa-va-drivers \
    mesa-vdpau-drivers \
    mesa-opencl-icd \
    ocl-icd-libopencl1 \
    vainfo 2>/dev/null || msg_warn "Some AMD APU packages failed"

  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"
    $STD apt install -y rocm-opencl-runtime rocm-smi-lib 2>/dev/null || msg_warn "ROCm minimal install also failed"
  }

  # 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
# ══════════════════════════════════════════════════════════════════════════════
_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"
    $STD apt-get -y install va-driver-all vainfo 2>/dev/null || true
    return 0
  fi

  msg_info "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"

    # For Debian 13 Trixie/Sid: Use Debian's own nvidia packages first (better compatibility)
    # NVIDIA's CUDA repo targets Debian 12 and may not have amd64 packages for Trixie
    if [[ "$os_codename" == "trixie" || "$os_codename" == "sid" ]]; then
      msg_info "Debian ${os_codename}: Using Debian's NVIDIA packages"

      # Extract major version for flexible matching (580.126.09 -> 580)
      local nvidia_major_version="${nvidia_host_version%%.*}"

      # Check what versions are actually 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 "Found available NVIDIA 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
          msg_ok "Installed NVIDIA libraries (${available_version})"
        else
          msg_warn "Failed to install NVIDIA ${available_version} - trying unversioned"
          $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends libcuda1 libnvcuvid1 libnvidia-encode1 libnvidia-ml1 2>/dev/null || true
        fi
      else
        # No matching major version - try latest available or unversioned
        msg_warn "No NVIDIA packages for version ${nvidia_major_version}.x found in repos"
        available_version=$(apt-cache madison libcuda1 2>/dev/null | awk '{print $3}' | head -1 || true)
        if [[ -n "$available_version" ]]; then
          msg_info "Trying latest available: ${available_version} (may cause version mismatch)"
          $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends \
            libcuda1="${available_version}" libnvcuvid1="${available_version}" \
            libnvidia-encode1="${available_version}" libnvidia-ml1="${available_version}" 2>/dev/null ||
            $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends \
              libcuda1 libnvcuvid1 libnvidia-encode1 libnvidia-ml1 2>/dev/null ||
            msg_warn "NVIDIA library installation failed - GPU compute may not work"
        else
          msg_warn "No NVIDIA packages available in Debian repos - GPU support disabled"
        fi
      fi
      $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends nvidia-smi 2>/dev/null || true

    else
      # Debian 11/12: Use NVIDIA CUDA repository for version matching
      local cuda_repo="debian12"
      case "$os_codename" in
      bullseye) cuda_repo="debian11" ;;
      bookworm) cuda_repo="debian12" ;;
      esac

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

      # Pin NVIDIA repo for version matching
      cat <<'NVIDIA_PIN' >/etc/apt/preferences.d/nvidia-cuda-pin
Package: *
Pin: origin developer.download.nvidia.com
Pin-Priority: 1001
NVIDIA_PIN

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

      # Extract major version for flexible matching (580.126.09 -> 580)
      local nvidia_major_version="${nvidia_host_version%%.*}"

      # Check what versions are actually available in CUDA repo
      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
          msg_ok "Installed version-matched NVIDIA libraries"
        else
          msg_warn "Version-pinned install failed - trying unpinned"
          $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends libcuda1 libnvcuvid1 libnvidia-encode1 libnvidia-ml1 2>/dev/null ||
            msg_warn "NVIDIA library installation failed"
        fi
      else
        msg_warn "No NVIDIA packages for version ${nvidia_major_version}.x in CUDA repo (host: ${nvidia_host_version})"
        # Try latest available version
        available_version=$(apt-cache madison libcuda1 2>/dev/null | awk '{print $3}' | head -1 || true)
        if [[ -n "$available_version" ]]; then
          msg_info "Trying latest available: ${available_version} (version mismatch warning)"
          if $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends \
            libcuda1="${available_version}" libnvcuvid1="${available_version}" \
            libnvidia-encode1="${available_version}" libnvidia-ml1="${available_version}" 2>/dev/null; then
            msg_ok "Installed NVIDIA libraries (${available_version}) - version differs from host"
          else
            $STD apt-get -y -o Dpkg::Options::="--force-confold" install --no-install-recommends libcuda1 libnvcuvid1 libnvidia-encode1 libnvidia-ml1 2>/dev/null ||
              msg_warn "NVIDIA library installation failed"
          fi
        else
          msg_warn "No NVIDIA packages available in CUDA repo - GPU support disabled"
        fi
      fi

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

  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
      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
        msg_ok "Installed version-matched NVIDIA libraries"
      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)
  $STD apt-get -y install va-driver-all vainfo 2>/dev/null || true

  # 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_info "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

  msg_ok "NVIDIA GPU configured"
}

# ══════════════════════════════════════════════════════════════════════════════
# 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_info "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

    # Verify nvidia-smi
    if command -v nvidia-smi &>/dev/null; then
      if nvidia-smi &>/dev/null; then
        msg_info "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_info "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
}