#!/bin/bash

# BigLinux GRUB Restore - System Detection Script
# www.biglinux.com.br
# Adapted for GTK4 frontend
# Updated and modified for Manjaro by Stefano Capitani, 2026
# License: GPL v3 or later

# Translation setup
export TEXTDOMAINDIR="/usr/share/locale"
export TEXTDOMAIN=grub-restore

# Function to log messages
log_message() {
    local msg="$(date '+%Y-%m-%d %H:%M:%S'): $1"
    echo "$msg" >&2
}

# Function to safely unmount partitions
safe_unmount() {
    local partition="$1"
    if mountpoint -q "$partition" 2>/dev/null; then
        log_message "Unmounting $partition"
        umount -l "$partition" 2>/dev/null || true
    fi
}

# Function to detect separate partitions from target system's fstab
detect_separate_partitions() {
    local root_partition="$1"
    local fstype="$2"
    local mount_point="/mnt/detect_$$"
    
    log_message "Detecting separate partitions for $root_partition (fstype: $fstype)..."
    
    # Create temporary mount point
    mkdir -p "$mount_point"
    
    # Mount root partition
    local mount_opts=""
    if [ "$fstype" = "btrfs" ]; then
        mount_opts="-o subvol=@"
    fi
    
    if ! mount $mount_opts "$root_partition" "$mount_point" 2>/dev/null; then
        log_message "Failed to mount $root_partition for fstab detection"
        rmdir "$mount_point"
        return 1
    fi
    
    local fstab_file="$mount_point/etc/fstab"
    local separate_partitions_file="/tmp/separate-partitions-$(basename "$root_partition")"
    
    # Clean previous file
    rm -f "$separate_partitions_file"
    touch "$separate_partitions_file"
    chmod 666 "$separate_partitions_file"
    
    if [ -f "$fstab_file" ]; then
        log_message "Reading fstab from $root_partition"
        
        # Parse fstab for separate partitions
        # Format: device mountpoint fstype options dump pass
        while IFS= read -r line; do
            # Skip comments and empty lines
            [[ "$line" =~ ^[[:space:]]*# ]] && continue
            [[ -z "${line// }" ]] && continue
            
            # fstab fields
            local device=$(echo "$line" | awk '{print $1}')
            local mountpoint=$(echo "$line" | awk '{print $2}')
            local fstype=$(echo "$line" | awk '{print $3}')
            local options=$(echo "$line" | awk '{print $4}')
            
            # Skip root, swap, and invalid entries
            [[ "$mountpoint" == "/" ]] && continue
            [[ "$mountpoint" == "swap" ]] && continue
            [[ "$mountpoint" == "none" ]] && continue
            [[ "$fstype" == "swap" ]] && continue
            [[ -z "$device" ]] && continue
            [[ -z "$mountpoint" ]] && continue
            
            # Resolve device (UUID, LABEL, PATH)
            local resolved_device=""
            if [[ "$device" == UUID=* ]]; then
                local uuid="${device#UUID=}"
                resolved_device=$(blkid -U "$uuid" 2>/dev/null)
            elif [[ "$device" == LABEL=* ]]; then
                local label="${device#LABEL=}"
                resolved_device=$(blkid -L "$label" 2>/dev/null)
            elif [[ "$device" == PARTUUID=* ]]; then
                local partuuid="${device#PARTUUID=}"
                resolved_device=$(blkid -t PARTUUID="$partuuid" -o device 2>/dev/null)
            elif [[ "$device" == PARTLABEL=* ]]; then
                local partlabel="${device#PARTLABEL=}"
                resolved_device=$(blkid -t PARTLABEL="$partlabel" -o device 2>/dev/null)
            elif [[ "$device" == /dev/* ]]; then
                resolved_device="$device"
            fi
            
            if [ -n "$resolved_device" ] && [ -b "$resolved_device" ]; then
                # Verify it's not the root partition
                if [ "$(realpath "$resolved_device" 2>/dev/null)" != "$(realpath "$root_partition" 2>/dev/null)" ]; then
                    # Get UUID and label for resolved device
                    local part_uuid=$(blkid -s UUID -o value "$resolved_device" 2>/dev/null)
                    local part_label=$(blkid -s LABEL -o value "$resolved_device" 2>/dev/null)
                    local part_fstype=$(blkid -s TYPE -o value "$resolved_device" 2>/dev/null)
                    
                    # Use fstype from fstab if blkid doesn't find it
                    [ -z "$part_fstype" ] && part_fstype="$fstype"
                    
                    # Format: device:mountpoint:fstype:options:uuid:label
                    echo "$resolved_device:$mountpoint:$part_fstype:$options:$part_uuid:$part_label" >> "$separate_partitions_file"
                    log_message "Found separate partition: $resolved_device -> $mountpoint ($part_fstype)"
                fi
            fi
        done < "$fstab_file"
    else
        log_message "No fstab found in $root_partition"
    fi
    
    # Unmount
    umount -l "$mount_point" 2>/dev/null || true
    rmdir "$mount_point" 2>/dev/null || true
    
    # Return file path
    echo "$separate_partitions_file"
}

# Garantir que o script rode como root
if [[ $EUID -ne 0 ]]; then
   log_message "This script must be run as root."
   exit 1
fi

log_message "Starting system detection..."

# Initial cleanup and create temp files with correct permissions
rm -f /tmp/os-prober /tmp/efi-partitions /tmp/grub-disks /tmp/separate-partitions-*
touch /tmp/os-prober /tmp/efi-partitions /tmp/grub-disks
chmod 666 /tmp/os-prober /tmp/efi-partitions /tmp/grub-disks

# 1. Unmount partitions to avoid conflicts with os-prober
log_message "Unmounting existing partitions..."
# Added support for more FS types and better blkid filter
for i in $(blkid -o device); do
    type=$(blkid -s TYPE -o value "$i")
    case "$type" in
        btrfs|ext[2-4]|xfs|vfat|ntfs)
            safe_unmount "$i"
            ;;
    esac
done

# 2. Run os-prober to detect Linux systems
log_message "Running os-prober to detect Linux systems..."
# os-prober needs partitions unmounted or manages them itself
GRUB_DISABLE_OS_PROBER=false os-prober | grep :linux | grep -v "timeshift" | grep -v 'subvol=@/' > /tmp/os-prober 2>/dev/null

# Post-processing to ensure UUID and FSType are present
# Expected format by Python backend: part:label:desc:type:fstype:uuid
# Default os-prober often returns only: part:label:desc:type
if [ -s /tmp/os-prober ]; then
    log_message "Post-processing os-prober results to add missing UUIDs..."
    temp_file=$(mktemp)
    while IFS= read -r line; do
        if [ -z "$line" ]; then continue; fi
        
        # Check if already has UUID (simple check)
        if [[ "$line" != *"UUID="* ]]; then
            part=$(echo "$line" | cut -d: -f1)
            
            uuid=$(blkid -s UUID -o value "$part")
            fstype=$(blkid -s TYPE -o value "$part")
            
            # Remove trailing colons
            clean_line=$(echo "$line" | sed 's/:*$//')
            
            # Add missing fields.
            # Assume original line has 4 fields. Adjust if less.
            # Target format: part:label:desc:type:fstype:UUID=...
            
            # Count current fields
            num_fields=$(echo "$clean_line" | awk -F: '{print NF}')
            
            # Rebuild line for consistent structure
            # If 4 fields, add fstype and uuid
            if [ "$num_fields" -le 4 ]; then
                new_line="${clean_line}:${fstype}:UUID=${uuid}"
                echo "$new_line" >> "$temp_file"
            else
                # If already more fields but no UUID, try to append
                new_line="${clean_line}:UUID=${uuid}"
                echo "$new_line" >> "$temp_file"
            fi
        else
            echo "$line" >> "$temp_file"
        fi
    done < /tmp/os-prober
    mv "$temp_file" /tmp/os-prober
    chmod 666 /tmp/os-prober
fi

# Add extra Btrfs partitions (important for specific subvolumes)
log_message "Checking for additional btrfs partitions..."
for i in $(blkid -t TYPE="btrfs" -o device); do
    if ! grep -q "$i" /tmp/os-prober 2>/dev/null; then
        label=$(blkid -s LABEL -o value "$i")
        if [ "$label" != "timeshift-btrfs" ]; then
            log_message "Adding btrfs partition: $i"
            [ -z "$label" ] && label="Btrfs_System"
            
            # Get UUID for consistency
            uuid=$(blkid -s UUID -o value "$i")
            
            # Adjusted format for Python parser compatibility
            # part:label:desc:type:fstype:UUID=uuid
            # Leave fields 3 and 4 empty (os-prober desc and type)
            echo "$i:$label:Linux::btrfs:UUID=$uuid" >> /tmp/os-prober 
        fi
    fi
done

# 3. Detect separate partitions for each detected system
log_message "Detecting separate partitions for each system..."
if [ -s /tmp/os-prober ]; then
    while IFS= read -r line; do
        if [ -z "$line" ]; then continue; fi
        part=$(echo "$line" | cut -d: -f1)
        fstype=$(echo "$line" | cut -d: -f5)
        # If fstype not in field 5, try field 4
        if [ -z "$fstype" ] || [ "$fstype" = "UUID="* ]; then
            fstype=$(echo "$line" | cut -d: -f4)
        fi
        # Default to ext4 if not determined
        [ -z "$fstype" ] && fstype="ext4"
        
        log_message "Checking separate partitions for $part ($fstype)"
        detect_separate_partitions "$part" "$fstype"
    done < /tmp/os-prober
fi

# 4. EFI Detection

# Improved live boot device detection (avoid formatting the USB stick by mistake)
DEVICE_LIVE_BOOT=$(lsblk -no pkname $(findmnt -n -o SOURCE /run/miso/bootmnt 2>/dev/null) 2>/dev/null)
log_message "Live boot device to exclude: $DEVICE_LIVE_BOOT"

# Scan for EFI partitions (Force Mode - All VFAT partitions)
# Using lsblk to be more robust and list all VFAT/FAT32 without mounting checks
log_message "Scanning for potential EFI (vfat/fat) partitions..."
lsblk -ln -o PATH,FSTYPE | awk 'tolower($2) ~ /fat/ {print $1}' | while read -r i; do
    # Skip if it's the live boot device
    if [ -n "$DEVICE_LIVE_BOOT" ]; then
        [[ "$i" == *"$DEVICE_LIVE_BOOT"* ]] && continue
    fi
    
    log_message "Found potential EFI partition: $i"
    echo "$i" >> /tmp/efi-partitions
done

# 5. Physical Disk Detection (Improved for NVMe and SSDs)
log_message "Starting disk detection..."
# lsblk -d (physical devices only) -n (no header)
for dev_raw in $(lsblk -dno NAME | grep -v "zram" | grep -v "loop" | grep -v "sr"); do
    # Normalize dev name (strip /dev/ prefix if present)
    dev=${dev_raw#/dev/}

    # Skip live boot disk (only if defined)
    if [ -n "$DEVICE_LIVE_BOOT" ]; then
        clean_live_boot=${DEVICE_LIVE_BOOT#/dev/}
        if [[ "$dev" == "$clean_live_boot" ]]; then
            log_message "Skipping live boot device (exact match): $dev"
            continue
        fi
    fi
    
    # Get details individually to avoid parsing issues
    size=$(lsblk -dn -o SIZE "/dev/$dev" | head -n1)
    if [ -z "$size" ]; then size="Unknown"; fi
    
    model=$(lsblk -dn -o MODEL "/dev/$dev" | head -n1)
    # Basic string cleanup
    clean_model=$(echo "${model:-Unknown_Disk}" | tr ' ' '_')
    
    pttype=$(lsblk -dn -o PTTYPE "/dev/$dev" | head -n1)
    clean_pttype=$(echo "${pttype:-unknown}" | tr ' ' '_')
    
    echo "/dev/${dev} ${size} ${clean_model} ${clean_pttype}" >> /tmp/grub-disks
    log_message "Added disk: /dev/${dev} (${size}, ${clean_model}, ${clean_pttype})"
done

log_message "Detection completed. Systems: $(wc -l < /tmp/os-prober), EFI: $(wc -l < /tmp/efi-partitions), Disks: $(wc -l < /tmp/grub-disks)"

exit 0
