#!/bin/bash
#
# modules/user/nvidia.sh — install the NVIDIA driver branch that matches the
# GPU that is actually in the machine.
#
# Runs on its own:
#     bash modules/user/nvidia.sh
#     LSPCI_BIN=/tmp/fake-lspci bash modules/user/nvidia.sh
#
# Behaviour:
#   * no NVIDIA GPU            -> prints one line and exits 0
#   * GPU found, SB off/setup  -> installs, no signing, no MOK enrolment
#   * GPU found, SB enabled    -> generates + imports the akmods key, then STOPS
#                                 and asks for a reboot. The driver is installed
#                                 on the next run, once the key is enrolled.
#                                 (Installing before enrolment is the documented
#                                 way to get a black screen: RPM Fusion blacklists
#                                 nouveau the moment the driver package lands.)
#   * already installed        -> skips the transaction, keeps `dnf mark user`
#
set -uo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=../../lib/common.sh
source "$HERE/../../lib/common.sh"
# shellcheck source=../../lib/ui.sh
source "$HERE/../../lib/ui.sh"

if is_debug; then set +e; else set -e; fi

# Overridable so the detection logic can be tested on a host with no NVIDIA
# hardware (same pattern as OS_RELEASE_FILE / FSTAB_FILE in lib/common.sh).
LSPCI_BIN="${LSPCI_BIN:-lspci}"
MOKUTIL_BIN="${MOKUTIL_BIN:-mokutil}"
EFIVARS_DIR="${EFIVARS_DIR:-/sys/firmware/efi/efivars}"
AKMODS_CERT="${AKMODS_CERT:-/etc/pki/akmods/certs/public_key.der}"
AKMOD_COMPILE_TIMEOUT="${AKMOD_COMPILE_TIMEOUT:-300}"

step "NVIDIA driver"
rule

#--------------------------------------------------
# Detection
#--------------------------------------------------
# Prints: "<branch>|<model>|<pciid>"  or nothing at all when there is no card.
#
# Branch is derived from the chip family in the name lspci prints (hwdata always
# has it):  GB/AD/GA/TU -> current,  GV/GP/GM -> 580xx,  GK -> 470xx,  GF -> 390xx.
# The numeric device id is NOT used for the decision: the id ranges overlap
# between families, and a wrong guess installs a module the card cannot use.
detect_nvidia_gpu() {
    local line model arch id

    line=$("$LSPCI_BIN" -nn 2>/dev/null \
        | grep -iE '(VGA compatible|3D|Display) controller' \
        | grep -iE 'nvidia|\[10de:' \
        | head -n1) || true
    [[ -z "$line" ]] && return 1

    id=$(printf '%s' "$line" \
        | grep -oE '\[10de:[0-9a-fA-F]{4}\]' | head -n1 | tr -d '[]') || true

    # "01:00.0 VGA compatible controller [0300]: NVIDIA Corporation GA102
    #  [GeForce RTX 3080 Ti] [10de:2208] (rev a1)"  ->  "GA102 [GeForce RTX 3080 Ti]"
    model=$(printf '%s' "$line" \
        | sed -n 's/.*NVIDIA Corporation \(.*\)\[10de:[0-9a-fA-F]\{4\}\].*/\1/p' \
        | sed -e 's/[[:space:]]*$//') || true
    [[ -z "$model" ]] && model="unknown NVIDIA device"

    arch=$(printf '%s' "$model" | grep -oE '(GB|AD|GA|TU|GV|GP|GM|GK|GF)[0-9]{3}[A-Z]?' \
        | head -n1 | grep -oE '^(GB|AD|GA|TU|GV|GP|GM|GK|GF)') || true

    local branch="unknown"
    case "$arch" in
        GB|AD|GA|TU) branch="current" ;;   # Turing and newer: open kernel module
        GV|GP|GM)    branch="580xx"   ;;   # Volta / Pascal / Maxwell (last branch)
        GK)          branch="470xx"   ;;   # Kepler, Xorg only
        GF)          branch="390xx"   ;;   # Fermi, Xorg only, EOL upstream
        "")
            # Pre-Fermi naming (G80/G92/GT200/GT215...) has no architecture
            # prefix; those cards need the EOL 340xx branch plus a longterm
            # kernel, which this script deliberately does not automate.
            if printf '%s' "$model" | grep -qE '(^|[^A-Za-z])(G[0-9]{2,3}|GT[0-9]{3}|NV[0-9]{2})'; then
                branch="340xx"
            else
                branch="unknown"
            fi
            ;;
    esac

    printf '%s|%s|%s\n' "$branch" "$model" "${id:-10de:unknown}"
}

gpu_info=""
if ! command -v "$LSPCI_BIN" >/dev/null 2>&1; then
    info "pciutils missing (no lspci). Installing..."
    install_optional pciutils
fi
if command -v "$LSPCI_BIN" >/dev/null 2>&1; then
    gpu_info=$(detect_nvidia_gpu) || true
fi

if [[ -z "$gpu_info" ]]; then
    info "No NVIDIA GPU detected. Nothing to do."
    exit 0
fi

BRANCH="${gpu_info%%|*}"
rest="${gpu_info#*|}"
MODEL="${rest%%|*}"
PCIID="${rest#*|}"

ok "GPU: $MODEL ($PCIID)"

# Multiple cards of different generations are not something one dnf transaction
# can serve correctly: the packages are mutually exclusive.
nv_count=$("$LSPCI_BIN" -nn 2>/dev/null \
    | grep -icE '(VGA compatible|3D|Display) controller.*(nvidia|\[10de:)') || true
if [[ "${nv_count:-0}" -gt 1 ]]; then
    warn "Found $nv_count NVIDIA display devices. Only the first one determines"
    warn "the driver branch; mixed generations on one machine need manual work."
fi

# A 3D controller (rather than a VGA/Display controller) is the Optimus layout.
if printf '%s' "$("$LSPCI_BIN" -nn 2>/dev/null | grep -i '10de:' | head -n1)" \
    | grep -q '3D controller'; then
    info "Hybrid (Optimus) layout: Fedora handles the iGPU/dGPU split automatically."
fi

#--------------------------------------------------
# Branch -> packages
#--------------------------------------------------
if [[ "$BRANCH" == "unknown" ]]; then
    warn "Could not identify the GPU generation from '$MODEL'."
    choice=$(ui_choose "NVIDIA driver branch" "Which driver branch does this card need?" \
        "current (Turing / Ampere / Ada / Blackwell)" \
        "current (Turing / Ampere / Ada / Blackwell)" \
        "580xx (Maxwell / Pascal / Volta)" \
        "470xx (Kepler)" \
        "390xx (Fermi)")
    case "$choice" in
        580xx*) BRANCH="580xx" ;;
        470xx*) BRANCH="470xx" ;;
        390xx*) BRANCH="390xx" ;;
        *)      BRANCH="current" ;;
    esac
fi

case "$BRANCH" in
    current)
        PKGS=(akmod-nvidia xorg-x11-drv-nvidia-cuda)
        MARK="akmod-nvidia"
        BRANCH_LABEL="current (Turing/Ampere/Ada/Blackwell, open kernel module)"
        ;;
    580xx)
        PKGS=(akmod-nvidia-580xx xorg-x11-drv-nvidia-580xx xorg-x11-drv-nvidia-580xx-cuda)
        MARK="akmod-nvidia-580xx"
        BRANCH_LABEL="580xx (Maxwell/Pascal/Volta, last branch for those cards)"
        ;;
    470xx)
        PKGS=(akmod-nvidia-470xx xorg-x11-drv-nvidia-470xx xorg-x11-drv-nvidia-470xx-cuda)
        MARK="akmod-nvidia-470xx"
        BRANCH_LABEL="470xx (Kepler, Xorg only)"
        ;;
    390xx)
        PKGS=(akmod-nvidia-390xx xorg-x11-drv-nvidia-390xx xorg-x11-drv-nvidia-390xx-cuda)
        MARK="akmod-nvidia-390xx"
        BRANCH_LABEL="390xx (Fermi, EOL upstream, Xorg only)"
        ;;
    340xx)
        warn "This card predates Fermi. Driver 340xx is end-of-life and needs a"
        warn "longterm kernel from a COPR, which this script does not automate."
        warn "See https://rpmfusion.org/Howto/NVIDIA and install it by hand."
        exit 0
        ;;
esac

info "Driver branch: $BRANCH_LABEL"
if [[ "$BRANCH" != "current" ]]; then
    warn "This branch has no Wayland support. Plan on an X11 session."
fi

#--------------------------------------------------
# Secure Boot gate
#--------------------------------------------------
# enabled -> sign and enrol, then stop for a reboot
# disabled / setup mode -> no signing at all (an import request would never be
#                          processed: shim skips MokManager when enforcement is
#                          off, so the request just sits in the efivars)
sb_state() {
    local out sb_var pk_var

    if command -v "$MOKUTIL_BIN" >/dev/null 2>&1; then
        out=$("$MOKUTIL_BIN" --sb-state 2>/dev/null) || true
        case "$out" in
            *"SecureBoot enabled"*)  echo "enabled";  return 0 ;;
            *"Setup mode"*)          echo "setup";    return 0 ;;
            *"SecureBoot disabled"*) echo "disabled"; return 0 ;;
        esac
    fi

    # Fallback: read the efivars directly (byte 5 of the variable data).
    if [[ -d "$EFIVARS_DIR" ]]; then
        sb_var=$(ls "$EFIVARS_DIR"/SecureBoot-* 2>/dev/null | head -n1) || true
        pk_var=$(ls "$EFIVARS_DIR"/PK-* 2>/dev/null | head -n1) || true
        if [[ -n "$sb_var" ]]; then
            if [[ "$(od -An -t u1 -j 4 -N 1 "$sb_var" 2>/dev/null | tr -d '[:space:]')" == "1" ]]; then
                echo "enabled"; return 0
            fi
            # No Platform Key enrolled = Setup mode: nothing can be enforced.
            [[ -z "$pk_var" ]] && { echo "setup"; return 0; }
            echo "disabled"; return 0
        fi
    fi

    echo "unknown"
}

SB_STATE=$(sb_state)
case "$SB_STATE" in
    enabled)
        info "Secure Boot is enabled: the module has to be signed and the key enrolled."
        install_packages akmods

        if [[ ! -f "$AKMODS_CERT" ]]; then
            info "Generating the akmods signing key..."
            sudo kmodgenca -a
        else
            ok "akmods signing key present."
        fi

        if sudo "$MOKUTIL_BIN" --test-key "$AKMODS_CERT" 2>&1 | grep -qi 'not enrolled'; then
            warn "The akmods key is not enrolled yet. Importing it now."
            sudo "$MOKUTIL_BIN" --import "$AKMODS_CERT"
            printf "\n"
            color_yellow "------------------------------------------------------------"
            color_yellow " ACTION REQUIRED: enrolment is pending"
            color_yellow "------------------------------------------------------------"
            info "The driver is deliberately NOT installed yet: with Secure Boot"
            info "enforcing, an unsigned module leaves you with a black screen."
            printf "\n"
            info "1. Reboot now."
            info "2. At the boot screen press a key during the countdown."
            info "3. Enroll MOK -> View key 0 (check the fingerprint) -> Continue"
            info "   -> Yes -> type the password you just set. The keyboard there"
            info "   is QWERTY, so use a short, symbol-free password."
            info "4. Boot Fedora and run this setup again: it will install the"
            info "   driver now that the key is trusted."
            printf "\n"
            warn "If you miss the countdown the request is dropped silently and"
            warn "you have to import again."
            exit 0
        fi
        ok "akmods key already enrolled."
        ;;
    setup)
        ok "Secure Boot is off (platform is in Setup mode): no signing needed."
        warn "Setup mode is a temporary state: a firmware update or 'Restore"
        warn "Factory Keys' turns enforcement on, and this unsigned module then"
        warn "stops loading at the next boot. Enrol the key before you need it:"
        warn "  sudo dnf install akmods && sudo kmodgenca -a"
        warn "  sudo mokutil --import /etc/pki/akmods/certs/public_key.der"
        ;;
    disabled)
        ok "Secure Boot is disabled: no signing needed."
        warn "Keep it that way deliberately: a firmware update or 'Restore Factory"
        warn "Keys' can enable enforcement, and the unsigned module then fails to"
        warn "load at the next boot."
        ;;
    *)
        warn "Could not determine the Secure Boot state. Assuming it is off."
        warn "If Secure Boot actually is on, the driver will not load: check"
        warn "'mokutil --sb-state' and enrol the akmods key first."
        ;;
esac

#--------------------------------------------------
# Install
#--------------------------------------------------
if pkg_installed_any "${PKGS[@]}"; then
    ok "NVIDIA driver packages already installed."
else
    info "Installing: ${PKGS[*]}"
    if ! install_packages "${PKGS[@]}"; then
        err "Driver installation failed for branch $BRANCH."
        err "Check that RPM Fusion nonfree is enabled: dnf repolist | grep rpmfusion"
        exit 1
    fi
fi

# An akmod is removable as an "unused dependency" unless it is marked user
# installed; `dnf autoremove` has taken the driver (and its compiler) away on
# people before.
if ! sudo dnf mark user "$MARK" >/dev/null 2>&1; then
    warn "Could not mark $MARK as user-installed (dnf autoremove may remove it)."
fi

#--------------------------------------------------
# Wait for the module build
#--------------------------------------------------
# dnf returns before the kmod exists. Rebooting in that window gives a boot with
# no driver at all, because nouveau is already blacklisted by the driver package.
NEWEST_KERNEL=$(rpm -q --qf '%{VERSION}-%{RELEASE}.%{ARCH}\n' kernel-core 2>/dev/null \
    | sort -V | tail -n1) || true

info "Waiting for the akmod build (up to ${AKMOD_COMPILE_TIMEOUT}s)..."
built=false
deadline=$(( SECONDS + AKMOD_COMPILE_TIMEOUT ))
while (( SECONDS < deadline )); do
    if [[ -n "$NEWEST_KERNEL" ]] && modinfo -k "$NEWEST_KERNEL" -F version nvidia >/dev/null 2>&1; then
        built=true; break
    fi
    if [[ -z "$NEWEST_KERNEL" ]] && modinfo -F version nvidia >/dev/null 2>&1; then
        built=true; break
    fi
    sleep 5
done

if [[ "$built" == "true" ]]; then
    if [[ -n "$NEWEST_KERNEL" ]]; then
        VER=$(modinfo -k "$NEWEST_KERNEL" -F version nvidia 2>/dev/null) || true
    else
        VER=$(modinfo -F version nvidia 2>/dev/null) || true
    fi
    ok "Kernel module built: nvidia ${VER:-unknown}"
    if [[ -n "$NEWEST_KERNEL" ]] && ! [[ "$NEWEST_KERNEL" == "$(uname -r)" ]]; then
        warn "Built for $NEWEST_KERNEL but you are running $(uname -r): reboot."
    fi
else
    warn "The module is not built yet after ${AKMOD_COMPILE_TIMEOUT}s."
    warn "Do NOT reboot until it is: the next boot would have no driver."
    warn "Check the build log, then retry:"
    warn "  ls -l /var/cache/akmods/nvidia/"
    warn "  sudo akmods --force"
fi

#--------------------------------------------------
# Optional extras
#--------------------------------------------------
printf "\n"
extras=$(ui_confirm "Install the optional graphics extras (VA-API decode, vulkan)?" "no")
if [[ "$extras" == "yes" ]]; then
    install_optional libva-nvidia-driver libva-utils vulkan
fi

#--------------------------------------------------
# What to check after the reboot
#--------------------------------------------------
printf "\n"
info "After the reboot, verify with:"
info "    nvidia-smi"
info "    cat /sys/module/nvidia_drm/parameters/modeset     # expect Y"
info "    lsmod | grep -Ei 'nouveau|nova_core'              # expect nothing"
warn "A reboot is required before the driver is active."