#!/bin/bash
# ==============================================================================
# Snarfnet Raspberry Pi SSD Boot Diagnostic & Automated Repair Script
# Run this script after booting an offline node (.21, .27, .73, .75) from a
# rescue SD card with the node's target SSD attached.
# ==============================================================================
set -euo pipefail

echo "=================================================================="
echo " Snarfnet RPi Offline Node SSD Diagnostic & Repair Tool"
echo "=================================================================="

if [ "$(id -u)" -ne 0 ]; then
    echo "[ERROR] Please run as root (sudo bash $0)"
    exit 1
fi

# 1. Identify which device is currently mounted as the Rescue OS root (/)
RESCUE_ROOT_PART=$(findmnt -n -o SOURCE /)
RESCUE_DISK=$(lsblk -no PKNAME "$RESCUE_ROOT_PART" 2>/dev/null | head -n1)
echo "[INFO] Rescue OS is running from: $RESCUE_ROOT_PART (disk: /dev/$RESCUE_DISK)"

echo ""
echo "--- Attached Block Devices ---"
lsblk -o NAME,SIZE,TYPE,FSTYPE,LABEL,MOUNTPOINT,PARTUUID,MODEL,TRAN
echo "------------------------------"

# 2. Auto-detect target SSD (prefer /dev/sda or /dev/nvme0n1 that is NOT the rescue disk)
TARGET_DISK="${1:-}"
if [ -z "$TARGET_DISK" ]; then
    for cand in /dev/sda /dev/nvme0n1 /dev/sdb; do
        if [ -b "$cand" ] && [ "$(basename "$cand")" != "$RESCUE_DISK" ]; then
            TARGET_DISK="$cand"
            break
        fi
    done
fi

if [ -z "$TARGET_DISK" ] || [ ! -b "$TARGET_DISK" ]; then
    echo "[WARN] Target SSD not detected yet!"
    echo " -> If this is a Shipyard Pi 5 (.73 or .75), plug in the USB SSD now and wait 5 seconds."
    echo " -> Usage: sudo bash $0 /dev/sda"
    exit 1
fi

echo "[INFO] Target drive selected for repair: $TARGET_DISK"

# Determine partition names (/dev/sda1 + /dev/sda2 vs /dev/nvme0n1p1 + /dev/nvme0n1p2)
if [[ "$TARGET_DISK" == *"nvme"* ]] || [[ "$TARGET_DISK" == *"mmcblk"* ]]; then
    BOOT_PART="${TARGET_DISK}p1"
    ROOT_PART="${TARGET_DISK}p2"
else
    BOOT_PART="${TARGET_DISK}1"
    ROOT_PART="${TARGET_DISK}2"
fi

if [ ! -b "$BOOT_PART" ] || [ ! -b "$ROOT_PART" ]; then
    echo "[ERROR] Expected partitions $BOOT_PART and $ROOT_PART not found on $TARGET_DISK!"
    lsblk "$TARGET_DISK"
    exit 1
fi

# Ensure target partitions are unmounted before running fsck
umount "$BOOT_PART" 2>/dev/null || true
umount "$ROOT_PART" 2>/dev/null || true

echo ""
echo "=== STEP 1: Checking SMART / USB Bridge Health ==="
lsusb || true
dmesg | grep -iE "usb|uas|sda|nvme|I/O error|reset" | tail -n 15 || true

echo ""
echo "=== STEP 2: Repairing Boot FAT32 Partition ($BOOT_PART) ==="
fsck.vfat -a -v "$BOOT_PART" || true

echo ""
echo "=== STEP 3: Repairing Root EXT4 Partition ($ROOT_PART) ==="
echo "Before repair superblock state:"
tune2fs -l "$ROOT_PART" | grep -E "Filesystem state|Mount count|Last mount time|Last write time|Errors behavior" || true

# Force full ext4 journal & inode/bitmap repair (-f -y)
e2fsck -f -y -v "$ROOT_PART" || true

# Set ext4 error behavior to 'continue' so minor block errors don't force read-only remount on boot
tune2fs -e continue "$ROOT_PART"

echo "After repair superblock state:"
tune2fs -l "$ROOT_PART" | grep -E "Filesystem state|Errors behavior" || true

echo ""
echo "=== STEP 4: Mounting Target OS & Auditing Boot Config ==="
mkdir -p /mnt/recovery
mount "$ROOT_PART" /mnt/recovery

if [ -d /mnt/recovery/boot/firmware ]; then
    BOOT_MNT="/mnt/recovery/boot/firmware"
else
    BOOT_MNT="/mnt/recovery/boot"
fi
mount "$BOOT_PART" "$BOOT_MNT"

NODE_HOSTNAME=$(cat /mnt/recovery/etc/hostname 2>/dev/null || echo "unknown")
echo "[INFO] Mounted Node Hostname: $NODE_HOSTNAME"

ACTUAL_BOOT_PARTUUID=$(blkid -s PARTUUID -o value "$BOOT_PART")
ACTUAL_ROOT_PARTUUID=$(blkid -s PARTUUID -o value "$ROOT_PART")
echo "[INFO] Actual PARTUUIDs -> Boot ($BOOT_PART): $ACTUAL_BOOT_PARTUUID | Root ($ROOT_PART): $ACTUAL_ROOT_PARTUUID"

echo ""
echo "--- Current /etc/fstab on $NODE_HOSTNAME ---"
cat /mnt/recovery/etc/fstab

# Backup fstab before hardening
cp /mnt/recovery/etc/fstab /mnt/recovery/etc/fstab.bak.$(date +%s)

# Fix 1: Add 'nofail,x-systemd.automount' to GlusterFS and NFS network mounts so they NEVER hang boot
if grep -qE "glusterfs|nfs" /mnt/recovery/etc/fstab; then
    sed -i -E 's/(glusterfs[[:space:]]+defaults,_netdev)/\1,nofail,x-systemd.automount/g' /mnt/recovery/etc/fstab
    sed -i -E 's/(nfs[[:space:]]+defaults,nfsvers=4,_netdev)/\1,nofail,x-systemd.automount/g' /mnt/recovery/etc/fstab
    echo "[FIXED] Added 'nofail,x-systemd.automount' to network mounts in /etc/fstab."
fi

# Fix 2: Ensure cmdline.txt has 'fsck.repair=yes' and 'rootwait' (missing on DietPi nodes!)
for cmdfile in "$BOOT_MNT/cmdline.txt" "/mnt/recovery/boot/cmdline.txt"; do
    if [ -f "$cmdfile" ] && ! grep -q "DO NOT EDIT" "$cmdfile"; then
        cp "$cmdfile" "${cmdfile}.bak.$(date +%s)"
        if ! grep -q "fsck.repair=yes" "$cmdfile"; then
            sed -i 's/rootwait/fsck.repair=yes rootwait/g' "$cmdfile"
            echo "[FIXED] Added 'fsck.repair=yes' to $cmdfile"
        fi
        echo "Updated $cmdfile:"
        cat "$cmdfile"
    fi
done

echo ""
echo "=== STEP 5: Checking Previous Boot Logs on Target Drive ==="
if [ -d /mnt/recovery/var/log/journal ]; then
    journalctl -D /mnt/recovery/var/log/journal -b -1 -p 0..3 --no-pager | tail -n 30 || true
else
    echo "[INFO] Node uses RAMlog (/var/log on tmpfs); no persistent journal on disk."
fi

echo ""
echo "=== STEP 6: Clean Unmount ==="
sync
umount "$BOOT_MNT"
umount /mnt/recovery
sync

echo "=================================================================="
echo " [SUCCESS] Repair completed for $NODE_HOSTNAME ($TARGET_DISK)!"
echo " You can now run 'sudo poweroff', remove the rescue SD card,"
echo " and power the node back on to boot normally from its SSD."
echo "=================================================================="