#!/bin/bash
# =============================================================================
# Step 5: Install Patched amdgpu Module
# =============================================================================
# Installs the patched amdgpu.ko.zst module, backs up the original,
# runs depmod, and rebuilds initramfs.
#
# Usage: sudo bash 05_install_module.sh [path_to_amdgpu.ko.zst]
#   If no path given, looks in the default build location.
#
# IMPORTANT: Reboot after running this script!
# =============================================================================

set -euo pipefail

KERNEL_VER="$(uname -r)"
MODULE_DIR="/lib/modules/$KERNEL_VER/kernel/drivers/gpu/drm/amd/amdgpu"
ORIGINAL="$MODULE_DIR/amdgpu.ko.zst"

# Find the module to install
if [[ -n "${1:-}" ]]; then
    NEW_MODULE="$1"
else
    # Try default build location
    NEW_MODULE="$HOME/kernel-build/linux-cachyos/src/linux-*/drivers/gpu/drm/amd/amdgpu/amdgpu.ko.zst"
    NEW_MODULE=$(ls $NEW_MODULE 2>/dev/null | head -1 || true)
fi

if [[ -z "$NEW_MODULE" || ! -f "$NEW_MODULE" ]]; then
    echo "ERROR: No module found. Provide path as argument:"
    echo "  sudo bash $0 /path/to/amdgpu.ko.zst"
    exit 1
fi

echo "============================================"
echo " Installing Patched amdgpu Module"
echo "============================================"
echo ""
echo "  Kernel:  $KERNEL_VER"
echo "  Source:  $NEW_MODULE"
echo "  Target:  $ORIGINAL"
echo ""

# ─────────────────────────────────────────────
# 1. Backup Original Module
# ─────────────────────────────────────────────
echo "[1/4] Backing up original module..."
if [[ -f "$ORIGINAL" && ! -f "$ORIGINAL.orig" ]]; then
    sudo cp "$ORIGINAL" "$ORIGINAL.orig"
    echo "  Backup: $ORIGINAL.orig"
elif [[ -f "$ORIGINAL.orig" ]]; then
    echo "  Backup already exists: $ORIGINAL.orig"
else
    echo "  WARNING: Original module not found at $ORIGINAL"
fi

# ─────────────────────────────────────────────
# 2. Install New Module
# ─────────────────────────────────────────────
echo "[2/4] Installing patched module..."
sudo cp "$NEW_MODULE" "$ORIGINAL"
echo "  Installed: $(ls -lh "$ORIGINAL" | awk '{print $5, $9}')"

# ─────────────────────────────────────────────
# 3. Run depmod
# ─────────────────────────────────────────────
echo "[3/4] Running depmod..."
sudo depmod -a "$KERNEL_VER"
echo "  Done."

# ─────────────────────────────────────────────
# 4. Rebuild Initramfs
# ─────────────────────────────────────────────
echo "[4/4] Rebuilding initramfs..."

if command -v mkinitcpio &>/dev/null; then
    # Arch/CachyOS
    sudo mkinitcpio -P
    echo "  mkinitcpio completed."
elif command -v dracut &>/dev/null; then
    sudo dracut --force
    echo "  dracut completed."
elif command -v update-initramfs &>/dev/null; then
    sudo update-initramfs -u
    echo "  update-initramfs completed."
fi

# Update bootloader if Limine
if command -v limine-update &>/dev/null; then
    echo "  Updating Limine bootloader..."
    sudo limine-update
    echo "  Limine updated."
fi

echo ""
echo "============================================"
echo " Module installed successfully!"
echo "============================================"
echo ""
echo " REBOOT NOW to load the patched module:"
echo "   sudo reboot"
echo ""
echo " After reboot, run 06_verify_installation.sh"
echo ""
echo " To restore original module if needed:"
echo "   sudo cp $ORIGINAL.orig $ORIGINAL"
echo "   sudo depmod -a $KERNEL_VER"
echo "   sudo mkinitcpio -P && sudo limine-update"
echo "   sudo reboot"