diff --git a/fix-ugreen-led-after-kernel-update.sh b/fix-ugreen-led-after-kernel-update.sh
index 3c10700..81d20b5 100644
--- a/fix-ugreen-led-after-kernel-update.sh
+++ b/fix-ugreen-led-after-kernel-update.sh
@@ -8,10 +8,21 @@ echo "Detected kernel: $KERNEL"
 
 MODULE_NAME="led-ugreen"
 MODULE_FILE="led-ugreen.ko"
+MODPATH="/lib/modules/$KERNEL"
+
+echo ""
+echo "=== Checking kernel headers ==="
+if [ ! -d "$MODPATH/build" ]; then
+    echo "❌ Kernel headers missing for $KERNEL"
+    echo "→ Installing pve-headers-$KERNEL"
+    apt update
+    apt install -y "pve-headers-$KERNEL"
+else
+    echo "✓ Kernel headers present"
+fi
 
 echo ""
 echo "=== Checking if module exists on disk ==="
-MODPATH="/lib/modules/$KERNEL"
 FOUND=$(find "$MODPATH" -name "$MODULE_FILE" 2>/dev/null || true)
 
 if [ -z "$FOUND" ]; then
@@ -28,8 +39,9 @@ dkms status ugreen-led/0.1 || true
 
 echo ""
 if [ -z "$FOUND" ]; then
-    echo "=== DKMS: Reinstalling ugreen-led for kernel $KERNEL ==="
+    echo "=== DKMS: Building and installing ugreen-led for kernel $KERNEL ==="
     dkms remove ugreen-led/0.1 -k "$KERNEL" --force 2>/dev/null || true
+    dkms build ugreen-led/0.1 -k "$KERNEL"
     dkms install ugreen-led/0.1 -k "$KERNEL"
 else
     echo "Skipping DKMS rebuild (module already present)"
@@ -41,32 +53,24 @@ depmod -a "$KERNEL"
 
 echo ""
 echo "=== Loading module ==="
-modprobe -v "$MODULE_NAME" || {
-    echo "❌ Could not load module even after DKMS rebuild!"
-    exit 1
-}
+modprobe -v "$MODULE_NAME"
 
 echo ""
 echo "=== Probing UGREEN LEDs ==="
-UGREEN_PROBE="./ugreen_leds_controller/scripts/ugreen-probe-leds"
-
-if [[ -x "$UGREEN_PROBE" ]]; then
-    $UGREEN_PROBE
-else
-    echo "⚠️ Probe script not found: $UGREEN_PROBE"
-fi
+UGREEN_PROBE="/root/ugreen_leds_controller/scripts/ugreen-probe-leds"
+[ -x "$UGREEN_PROBE" ] && "$UGREEN_PROBE" || echo "⚠️ Probe script not found"
 
 echo ""
 echo "=== Restarting services ==="
 systemctl daemon-reload
 systemctl start ugreen-netdevmon@enp2s0
-systemctl status --no-pager ugreen-netdevmon@enp2s0 || true
-
-echo ""
-echo "=== Fixing Disk I/O Monitor ==="
 systemctl start ugreen-diskiomon
 systemctl enable ugreen-diskiomon
+
+echo ""
+echo "=== Status ==="
+systemctl status --no-pager ugreen-netdevmon@enp2s0 || true
 systemctl status --no-pager ugreen-diskiomon || true
 
 echo ""
-echo "=== Done. ==="
+echo "=== Done ==="