#!/bin/bash
set -e

echo ">>> [gpu] Running INSIDE target system"

GPU_INFO="$(lspci -nn | grep -E 'VGA|3D|Display' || true)"
echo ">>> [gpu] Detected devices:"
echo "$GPU_INFO"

pacman -Syy --noconfirm || true

install_pkgs() {
  pacman -S --needed --noconfirm "$@" || true
}

install_pkgs mesa libglvnd vulkan-icd-loader

if echo "$GPU_INFO" | grep -qi "Intel"; then
  echo ">>> [gpu] Installing Intel stack..."
  install_pkgs mesa vulkan-intel intel-media-driver
fi

if echo "$GPU_INFO" | grep -qi "AMD\|ATI"; then
  echo ">>> [gpu] Installing AMD stack..."
  install_pkgs mesa lib32-mesa vulkan-radeon lib32-vulkan-radeon vulkan-icd-loader lib32-vulkan-icd-loader vulkan-tools mesa-utils
fi

if echo "$GPU_INFO" | grep -qi "NVIDIA"; then
  echo ">>> [gpu] NVIDIA detected -> driver kivalasztasa..."

  NV_IDS=$( { lspci -n -d 10de::0300; lspci -n -d 10de::0302; lspci -n -d 10de::0380; } | awk '{print $3}' | cut -d: -f2 | sort -u || true )

  driver_family() {
    local id="$1" dec
    [[ "$id" =~ ^[0-9a-fA-F]{4}$ ]] || { echo unsupported; return; }
    dec=$((16#$id))
    if [ "$dec" -ge $((0x1E00)) ]; then echo open; return; fi
    if [ "$dec" -ge $((0x1340)) ] && [ "$dec" -le $((0x13FF)) ]; then echo 580xx; return; fi
    if [ "$dec" -ge $((0x1400)) ] && [ "$dec" -le $((0x1DFF)) ]; then echo 580xx; return; fi
    if [ "$dec" -ge $((0x0FC0)) ] && [ "$dec" -le $((0x103F)) ]; then echo 470xx; return; fi
    if [ "$dec" -ge $((0x1180)) ] && [ "$dec" -le $((0x11FF)) ]; then echo 470xx; return; fi
    if [ "$dec" -ge $((0x1280)) ] && [ "$dec" -le $((0x12BF)) ]; then echo 470xx; return; fi
    echo unsupported
  }

  BESTFAM=unsupported
  BESTRANK=0
  for id in $NV_IDS; do
    fam=$(driver_family "$id")
    echo ">>> [gpu] NVIDIA ${id} -> ${fam}"
    rank=0
    case "$fam" in
      open) rank=4 ;;
      580xx) rank=3 ;;
      470xx) rank=2 ;;
    esac
    if [ "$rank" -gt "$BESTRANK" ]; then BESTRANK=$rank; BESTFAM=$fam; fi
  done

  case "$BESTFAM" in
    open)
      echo ">>> [gpu] Turing+ -> nvidia-open-dkms"
      install_pkgs dkms linux-cachyos-headers nvidia-open-dkms nvidia-utils nvidia-settings lib32-nvidia-utils
      ;;
    580xx)
      echo ">>> [gpu] Pascal/Maxwell/Volta -> nvidia-580xx-dkms (legacy)"
      install_pkgs dkms linux-cachyos-headers nvidia-580xx-dkms nvidia-580xx-utils nvidia-580xx-settings lib32-nvidia-580xx-utils
      ;;
    470xx)
      echo ">>> [gpu] Kepler -> nvidia-470xx-dkms (legacy, unsupported)"
      install_pkgs dkms linux-cachyos-headers nvidia-470xx-dkms nvidia-470xx-utils nvidia-470xx-settings lib32-nvidia-470xx-utils
      ;;
    *)
      echo ">>> [gpu] Nem tamogatott NVIDIA (Fermi vagy regebbi) -> nouveau marad"
      ;;
  esac

  if echo "$GPU_INFO" | grep -qi "Intel\|AMD\|ATI"; then
    echo ">>> [gpu] Hybrid laptop detected -> installing nvidia-prime"
    install_pkgs nvidia-prime
  fi

  NV_MODULE_OK=0
  if [ "$BESTFAM" != unsupported ] && modinfo -k "$(ls /usr/lib/modules | head -1)" nvidia >/dev/null 2>&1; then
    NV_MODULE_OK=1
  fi

  if [ "$NV_MODULE_OK" -eq 1 ]; then
    echo ">>> [gpu] NVIDIA modul elerheto -> beallitasok..."

    mkdir -p /etc/modules-load.d
    cat > /etc/modules-load.d/nvidia.conf <<'EOF'
nvidia
nvidia_modeset
nvidia_uvm
nvidia_drm
EOF

    mkdir -p /etc/modprobe.d
    cat > /etc/modprobe.d/nvidia-drm.conf <<'EOF'
options nvidia_drm modeset=1
EOF

    cat > /etc/modprobe.d/99-nvidia-blacklist.conf <<'EOF'
blacklist nouveau
options nouveau modeset=0
EOF

    mkdir -p /etc/mkinitcpio.conf.d
    cat > /etc/mkinitcpio.conf.d/raveos-nvidia.conf <<'EOF'
MODULES=(nvidia nvidia_modeset nvidia_uvm nvidia_drm)
EOF
  else
    echo ">>> [gpu] NVIDIA modul NEM erheto el -> nouveau marad (nincs blacklist)"
  fi

  if pacman -Q plasma-desktop >/dev/null 2>&1; then
    echo ">>> [gpu] NVIDIA + Plasma -> installing X11 session (plasma-x11-session)"
    install_pkgs plasma-x11-session
    mkdir -p /etc/sddm.conf.d
    cat > /etc/sddm.conf.d/10-raveos-nvidia-x11.conf <<'EOF'
[Users]
DefaultSession=plasmax11.desktop
EOF
  fi

  echo ">>> [gpu] Rebuilding initramfs..."
  mkinitcpio -P || true
fi

echo ">>> [gpu] Done."