#!/bin/bash
set -uo pipefail
export LC_ALL=C

MOD="${1:-}"
case "$MOD" in
    --terv|--telepit|--rendszeren) ;;
    *) echo "használat: raveos-gpu-drivers --terv|--telepit|--rendszeren" >&2; exit 2 ;;
esac

GPU_INFO="$(lspci -nn | grep -E 'VGA|3D|Display' || true)"
INTEL=0; AMD=0; NVIDIA=0; VIRTUALIS=0
echo "$GPU_INFO" | grep -q '\[8086:' && INTEL=1
echo "$GPU_INFO" | grep -q '\[1002:' && AMD=1
echo "$GPU_INFO" | grep -q '\[10de:' && NVIDIA=1
echo "$GPU_INFO" | grep -qE '\[(1af4|1234|15ad|1b36|80ee):' && VIRTUALIS=1

nvidia_csalad() {
    local id="$1" dec
    [[ "$id" =~ ^[0-9a-fA-F]{4}$ ]] || { echo nem_tamogatott; 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 nem_tamogatott
}

CSALAD=""
if [ "$NVIDIA" = 1 ]; then
    legjobb=0
    CSALAD=nem_tamogatott
    for id in $( { lspci -n -d 10de::0300; lspci -n -d 10de::0302; lspci -n -d 10de::0380; } | awk '{print $3}' | cut -d: -f2 | sort -u ); do
        cs=$(nvidia_csalad "$id")
        case "$cs" in open) r=4 ;; 580xx) r=3 ;; 470xx) r=2 ;; *) r=0 ;; esac
        if [ "$r" -gt "$legjobb" ]; then legjobb=$r; CSALAD=$cs; fi
    done
fi

FEJLECEK=()
for kernel in $(pacman -Qq 2>/dev/null | grep -E '^linux(-[a-z0-9]+)?$' | grep -v -E -- '-(firmware|api|headers|docs|tools)$'); do
    FEJLECEK+=("${kernel}-headers")
done
[ ${#FEJLECEK[@]} -gt 0 ] || FEJLECEK=(linux-cachyos-headers)

KOZOS=(mesa libglvnd vulkan-icd-loader)
INTEL_CSOMAGOK=()
AMD_CSOMAGOK=()
NVIDIA_CSOMAGOK=()
[ "$INTEL" = 1 ] && INTEL_CSOMAGOK=(lib32-mesa vulkan-intel lib32-vulkan-intel intel-media-driver)
[ "$AMD" = 1 ] && AMD_CSOMAGOK=(lib32-mesa vulkan-radeon lib32-vulkan-radeon lib32-vulkan-icd-loader vulkan-tools mesa-utils)
case "$CSALAD" in
    open) NVIDIA_CSOMAGOK=(dkms "${FEJLECEK[@]}" nvidia-open-dkms nvidia-utils nvidia-settings lib32-nvidia-utils) ;;
    580xx) NVIDIA_CSOMAGOK=(dkms "${FEJLECEK[@]}" nvidia-580xx-dkms nvidia-580xx-utils nvidia-580xx-settings lib32-nvidia-580xx-utils) ;;
    470xx) NVIDIA_CSOMAGOK=(dkms "${FEJLECEK[@]}" nvidia-470xx-dkms nvidia-470xx-utils nvidia-470xx-settings lib32-nvidia-470xx-utils) ;;
esac
if [ ${#NVIDIA_CSOMAGOK[@]} -gt 0 ] && { [ "$INTEL" = 1 ] || [ "$AMD" = 1 ]; }; then
    NVIDIA_CSOMAGOK+=(nvidia-prime)
fi
CSOMAGOK=()
for p in "${KOZOS[@]}" "${INTEL_CSOMAGOK[@]}" "${AMD_CSOMAGOK[@]}" "${NVIDIA_CSOMAGOK[@]}"; do
    [[ " ${CSOMAGOK[*]} " == *" $p "* ]] || CSOMAGOK+=("$p")
done

if [ "$MOD" = "--terv" ]; then
    echo "intel=$INTEL"
    echo "amd=$AMD"
    echo "nvidia=$NVIDIA"
    echo "virtualis=$VIRTUALIS"
    echo "nvidia_csalad=$CSALAD"
    echo "kozos=${KOZOS[*]}"
    echo "intel_csomagok=${INTEL_CSOMAGOK[*]}"
    echo "amd_csomagok=${AMD_CSOMAGOK[*]}"
    echo "nvidia_csomagok=${NVIDIA_CSOMAGOK[*]}"
    echo "fejlecek=${FEJLECEK[*]}"
    echo "csomagok=${CSOMAGOK[*]}"
    exit 0
fi

[ "$(id -u)" -eq 0 ] || { echo "a telepítés root jogot igényel" >&2; exit 2; }

echo ">>> [gpu] Felismert eszközök:"
echo "$GPU_INFO"
echo ">>> [gpu] Telepítendő: ${CSOMAGOK[*]}"
if ! pacman -S --needed --noconfirm --color never "${CSOMAGOK[@]}"; then
    [ "$MOD" = "--rendszeren" ] && exit 1
    echo ">>> [gpu] Figyelem: nem minden csomag települt."
fi

if [ ${#NVIDIA_CSOMAGOK[@]} -gt 0 ]; then
    if [ "$MOD" = "--rendszeren" ]; then
        celkernel="$(uname -r)"
    else
        celkernel="$(ls /usr/lib/modules | sort -V | tail -1)"
    fi
    if modinfo -k "$celkernel" nvidia >/dev/null 2>&1; then
        echo ">>> [gpu] NVIDIA modul elérhető, beállítások..."
        install -d /etc/modules-load.d /etc/modprobe.d /etc/mkinitcpio.conf.d
        printf '%s\n' nvidia nvidia_modeset nvidia_uvm nvidia_drm > /etc/modules-load.d/nvidia.conf
        echo "options nvidia_drm modeset=1" > /etc/modprobe.d/nvidia-drm.conf
        printf '%s\n' "blacklist nouveau" "options nouveau modeset=0" > /etc/modprobe.d/99-nvidia-blacklist.conf
        echo "MODULES+=(nvidia nvidia_modeset nvidia_uvm nvidia_drm)" > /etc/mkinitcpio.conf.d/raveos-nvidia.conf
        if [ "$MOD" = "--rendszeren" ]; then
            mkinitcpio -P || exit 1
        fi
    else
        echo ">>> [gpu] Az NVIDIA modul nem érhető el, a nouveau marad."
    fi
    if pacman -Q plasma-desktop >/dev/null 2>&1; then
        echo ">>> [gpu] NVIDIA + Plasma: X11 munkamenet"
        pacman -S --needed --noconfirm --color never plasma-x11-session || true
        install -d /etc/sddm.conf.d
        printf '%s\n' "[Users]" "DefaultSession=plasmax11.desktop" > /etc/sddm.conf.d/10-raveos-nvidia-x11.conf
    fi
fi
echo ">>> [gpu] Kész."