#!/bin/bash
set -euo pipefail

# NVIDIA Kernel Module Recovery Script
# Purpose: Reinstall nvidia-580xx-dkms and rebuild kernel module after package cleanup removed it
# Target: mw-pfeddersheim-workstation (Manjaro, NVIDIA GPU)

MODE="apply"
if [[ "${1:-}" == "--check" ]]; then
  MODE="check"
  echo "=== Running in CHECK mode (diagnostic only, no changes) ==="
elif [[ "${1:-}" == "--help" ]]; then
  echo "Usage: $0 [--check | --apply | --help]"
  echo "  --check  Diagnose and show what would be done"
  echo "  --apply  Reinstall nvidia-580xx-dkms and rebuild module (requires sudo)"
  echo "  --help   Show this help"
  exit 0
fi

mkdir -p "$HOME/logs" 2>/dev/null || true
LOCAL_LOG="$HOME/logs/nvidia-recovery-$(date +%Y%m%d-%H%M%S).log"
BOLD=$(tput bold 2>/dev/null || true)
RED=$(tput setaf 1 2>/dev/null || true)
GREEN=$(tput setaf 2 2>/dev/null || true)
YELLOW=$(tput setaf 3 2>/dev/null || true)
RESET=$(tput sgr0 2>/dev/null || true)

log() { echo -e "$*" | tee -a "$LOCAL_LOG"; }
section() { log ""; log "${BOLD}--- $1 ---${RESET}"; }

echo "" | tee -a "$LOCAL_LOG"
log "============================================="
log "  NVIDIA Kernel Module Recovery"
log "  Date: $(date)"
log "  Log:  $LOCAL_LOG"
log "============================================="

# ============================================================
# 1. DIAGNOSTICS
# ============================================================
section "Diagnostic"

# Get current kernel
RUNNING_KERNEL=$(uname -r)
log "Running kernel: $RED$RUNNING_KERNEL$RESET"

KERNEL_VERSION="${RUNNING_KERNEL%-MANJARO}"  # 6.18.33-1
KERNEL_MAJOR="${KERNEL_VERSION%%.*}"          # 6
KERNEL_MINOR="${KERNEL_VERSION#*.}"           # 18.33-1
KERNEL_MINOR="${KERNEL_MINOR%%.*}"           # 18
KERNEL_PKG="linux${KERNEL_MAJOR}${KERNEL_MINOR}"  # linux618
log "Kernel package: $KERNEL_PKG"

# Check if nvidia module exists
MODULE_PATH="/usr/lib/modules/$RUNNING_KERNEL/extra/modules/nvidia.ko.zst"
MODULE_PATH2="/usr/lib/modules/$RUNNING_KERNEL/extra/nvidia.ko.zst"
MODULE_PATH3="/usr/lib/modules/$RUNNING_KERNEL/updates/dkms/nvidia.ko.zst"
MODULE_PATH4="/usr/lib/modules/$RUNNING_KERNEL/updates/nvidia.ko.zst"
MODULE_PATH5="/usr/lib/modules/$RUNNING_KERNEL/kernel/drivers/gpu/drm/nvidia/nvidia.ko.zst"

MODULE_FOUND=false
if [[ -f "$MODULE_PATH" ]] || [[ -f "$MODULE_PATH2" ]] || [[ -f "$MODULE_PATH3" ]] || [[ -f "$MODULE_PATH4" ]] || [[ -f "$MODULE_PATH5" ]]; then
  MODULE_FOUND=true
  log "${GREEN}NVIDIA module found at $MODULE_PATH3$RESET"
else
  log "${RED}NVIDIA module MISSING for kernel $RUNNING_KERNEL$RESET"
fi

# Check nvidia-utils
if command -v nvidia-smi &>/dev/null; then
  log "${GREEN}nvidia-smi: available$RESET"
else
  log "${RED}nvidia-smi: not found$RESET"
fi

# Check DKMS
if command -v dkms &>/dev/null; then
  log "${GREEN}dkms: available$RESET"
else
  log "${RED}dkms: not found$RESET"
fi

# Check if nvidia-580xx-dkms is installed
if pacman -Q nvidia-580xx-dkms &>/dev/null; then
  DKMS_INSTALLED=true
  log "${GREEN}nvidia-580xx-dkms: installed$(pacman -Qi nvidia-580xx-dkms 2>/dev/null | grep Version)"
else
  DKMS_INSTALLED=false
  log "${RED}nvidia-580xx-dkms: NOT installed$RESET"
fi

# Check if nvidia-utils is installed
if pacman -Q nvidia-utils &>/dev/null; then
  log "${GREEN}nvidia-utils: installed"
else
  log "${RED}nvidia-utils: NOT installed$RESET"
fi

# Check modules-load config
if [[ -f /etc/modules-load.d/mhwd-gpu.conf ]]; then
  log "modules-load config: /etc/modules-load.d/mhwd-gpu.conf exists"
  log "  Entries: $(grep -v '^#\|^$' /etc/modules-load.d/mhwd-gpu.conf | tr '\n' ' ')"
else
  log "${RED}No mhwd-gpu.conf found$RESET"
fi

# ============================================================
# 2. CHECK DEPENDENCIES
# ============================================================
section "Dependency Check"

MISSING=()
if ! command -v dkms &>/dev/null; then
  MISSING+=(dkms)
fi
if ! command -v mkinitcpio &>/dev/null; then
  MISSING+=(mkinitcpio)
fi

if [[ ${#MISSING[@]} -gt 0 ]]; then
  log "${RED}Missing base tools: ${MISSING[*]}$RESET"
  log "These will be installed as part of the fix."
fi

# Check available disk space
DISK_AVAIL=$(df -h / | awk 'NR==2 {print $4}')
log "Disk available: $DISK_AVAIL"

# ============================================================
# 3. APPLY FIX
# ============================================================
section "Applying Fix"

if [[ "$MODE" == "check" ]]; then
  log "${BOLD}DRY-RUN MODE${RESET}"
  log "Would execute:"
  log "  sudo pacman -Rns --noconfirm nvidia-utils  # remove conflicting non-580xx driver"
  log "  sudo pacman -S --needed nvidia-580xx-dkms nvidia-580xx-utils"
  log "  sudo pacman -S --needed linux618-headers   # kernel headers for DKMS"
  log "  sudo dkms autoinstall"
  log "  sudo mkinitcpio -P"
  log "  sudo reboot"
  exit 0
fi

# Confirm before destructive action
log "${BOLD}This will reinstall the NVIDIA kernel module driver.$RESET"
log "The system will reboot automatically."
read -rp "Continue? [y/N] " confirm
if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
  log "ABORTED by user."
  exit 0
fi

# Step 1: Resolve driver conflict, install 580xx suite
log "[1/5] Resolving NVIDIA driver conflict and installing 580xx suite..."
# Check if newer nvidia-utils conflicts with 580xx-dkms
if pacman -Q nvidia-utils &>/dev/null && ! pacman -Q nvidia-580xx-utils &>/dev/null; then
  log "Conflict: nvidia-utils (non-580xx) installed alongside nvidia-580xx-dkms."
  log "Removing conflicting nvidia-utils, installing nvidia-580xx-utils..."
  sudo pacman -Rns --noconfirm nvidia-utils 2>&1 | tee -a "$LOCAL_LOG"
  sudo pacman -S --needed --noconfirm nvidia-580xx-dkms nvidia-580xx-utils 2>&1 | tee -a "$LOCAL_LOG"
else
  sudo pacman -S --needed --noconfirm nvidia-580xx-dkms nvidia-580xx-utils 2>&1 | tee -a "$LOCAL_LOG"
fi

# Step 2: Ensure kernel headers are available for DKMS
log "[1b/4] Checking kernel headers for DKMS..."
KERNEL_HEADERS="${KERNEL_PKG}-headers"
if pacman -Q "$KERNEL_HEADERS" &>/dev/null; then
  log "${GREEN}Kernel headers ($KERNEL_HEADERS) already installed.$RESET"
else
  log "Installing $KERNEL_HEADERS for DKMS..."
  sudo pacman -S --needed --noconfirm "$KERNEL_HEADERS" 2>&1 | tee -a "$LOCAL_LOG"
fi

# Step 3: Rebuild DKMS module for current kernel
log "[2/5] Rebuilding DKMS module for kernel $RUNNING_KERNEL..."
sudo dkms autoinstall 2>&1 | tee -a "$LOCAL_LOG"

# Step 3: Verify module was created
log "[3/5] Verifying module..."
if [[ -f "$MODULE_PATH" ]] || [[ -f "$MODULE_PATH2" ]] || [[ -f "$MODULE_PATH3" ]] || [[ -f "$MODULE_PATH4" ]] || [[ -f "$MODULE_PATH5" ]]; then
  log "${GREEN}Module verified!$RESET"
  # Show module details
  modinfo_output=$(modinfo nvidia 2>/dev/null || echo "modinfo failed")
  log "  $modinfo_output"
else
  log "${RED}WARNING: Module file not found after DKMS build.$RESET"
  # Detect version for manual build
  DKMS_VERSION=$(pacman -Qi nvidia-580xx-dkms 2>/dev/null | grep Version | awk '{print $2}' | cut -d- -f1)
  [[ -z "$DKMS_VERSION" ]] && DKMS_VERSION="580.159.04"
  log "Attempting manual build (version: $DKMS_VERSION)..."
  sudo dkms build -m nvidia -v "$DKMS_VERSION" -k "$RUNNING_KERNEL" 2>&1 | tee -a "$LOCAL_LOG"
  sudo dkms install -m nvidia -v "$DKMS_VERSION" -k "$RUNNING_KERNEL" 2>&1 | tee -a "$LOCAL_LOG"
fi

# Step 4: Rebuild initramfs
log "[4/5] Rebuilding initramfs..."
sudo mkinitcpio -P 2>&1 | tee -a "$LOCAL_LOG"

# ============================================================
# 4. VERIFICATION
# ============================================================
section "Post-Recovery Verification"

log "Checking module loading..."
if sudo modprobe nvidia 2>&1; then
  log "${GREEN}nvidia module loaded successfully$RESET"
else
  log "${YELLOW}modprobe nvidia failed - module will be loaded at next boot$RESET"
fi

# Detect 580xx version for manual DKMS fallback
DKMS_VERSION=$(pacman -Qi nvidia-580xx-dkms 2>/dev/null | grep Version | awk '{print $2}' | cut -d- -f1 || echo "580.159.04")

# Manual DKMS fallback uses DKMS_VERSION variable

log ""
log "${BOLD}Recovery complete!${RESET}"
log ""
log "Before rebooting, verify:"
log "  1. Module exists: ls /usr/lib/modules/$(uname -r)/extra/nvidia*"
log "  2. nvidia-smi works: sudo nvidia-smi"
log ""
log "Then reboot:"
log "  sudo reboot"
log ""
log "Log saved to: $LOCAL_LOG"
log "============================================="