#!/usr/bin/env bash

# Stormux Gaming Image - USB to Disk Cloner
# Clones the entire USB system to an internal disk

set -e

error_exit() {
    echo "Error: $1"
    exit 1
}

# Function to find the source USB device
find_source_device() {
    # Look for the device we're currently running from
    local root_device
    root_device=$(findmnt -n -o SOURCE /)
    if [[ "$root_device" =~ ^/dev/(sd[a-z]|nvme[0-9]n[0-9]|mmcblk[0-9]) ]]; then
        # Extract just the device name (remove partition number)
        echo "$root_device" | sed 's/[0-9]*$//' | sed 's/p$//'
    else
        # Fallback: look for devices with STORMUX label
        local labeled_device
        labeled_device=$(lsblk -no NAME,LABEL | grep -i stormux | head -1 | awk '{print "/dev/" $1}' | sed 's/[0-9]*$//' | sed 's/p$//')
        if [[ -n "$labeled_device" ]]; then
            echo "$labeled_device"
        else
            return 1
        fi
    fi
}

# Function to detect target disks (excluding source)
detect_target_disks() {
    local source_device="$1"
    local disks=()
    
    while IFS= read -r disk; do
        # Skip if it's a partition, loop device, CD-ROM, or the source device
        # For partitions: skip sda1, nvme0n1p1, mmcblk0p1, but keep sda, nvme0n1, mmcblk0
        if [[ ! "$disk" =~ (sd[a-z][0-9]+|nvme[0-9]+n[0-9]+p[0-9]+|mmcblk[0-9]+p[0-9]+)$ ]] && [[ ! "$disk" =~ ^/dev/loop ]] && [[ ! "$disk" =~ ^/dev/sr ]] && [[ "$disk" != "$source_device" ]]; then
            if [[ -b "$disk" ]]; then
                disks+=("$disk")
            fi
        fi
    done < <(lsblk -dpno NAME 2>/dev/null)
    
    printf '%s\n' "${disks[@]}"
}

# Function to get disk info
get_disk_info() {
    local disk="$1"
    local size
    local model
    size=$(lsblk -dpno SIZE "$disk" 2>/dev/null | tr -d ' ')
    model=$(lsblk -dpno MODEL "$disk" 2>/dev/null | tr -d ' ' || echo "Unknown")
    echo "$size - $model"
}

# Function to regenerate UUIDs and rename partition labels to prevent boot conflicts
regenerate_partition_identifiers() {
    local target_device="$1"
    echo "Regenerating partition UUIDs and labels to prevent boot conflicts..."
    
    # Arrays to store old and new UUIDs for fstab update
    declare -A old_uuids
    declare -A new_uuids
    
    # Get all partitions on the target device
    local partitions=()
    while IFS= read -r partition; do
        if [[ "$partition" != "$target_device" && "$partition" =~ ^${target_device}.+ ]]; then
            partitions+=("$partition")
        fi
    done < <(lsblk -lpno NAME "$target_device" 2>/dev/null)
    
    # Process each partition
    for partition in "${partitions[@]}"; do
        local fstype
        local current_label
        local current_uuid
        fstype=$(lsblk -no FSTYPE "$partition" 2>/dev/null)
        current_label=$(lsblk -no LABEL "$partition" 2>/dev/null)
        current_uuid=$(lsblk -no UUID "$partition" 2>/dev/null)
        
        # Skip if no filesystem
        [[ -z "$fstype" ]] && continue
        
        # Store old UUID if it exists
        [[ -n "$current_uuid" ]] && old_uuids["$partition"]="$current_uuid"
        
        # Determine new label based on current label and add -HDD suffix
        local new_label=""
        if [[ -n "$current_label" ]]; then
            case "$current_label" in
                "STORMUX"|"stormux")
                    new_label="STORMUX-HDD"
                    ;;
                "BOOT"|"boot"|"ESP")
                    new_label="BOOT-HDD"
                    ;;
                *)
                    # For any other label, just add -HDD suffix, truncate if too long
                    new_label="${current_label}-HDD"
                    ;;
            esac
        fi
        
        # Handle the known 3-partition structure
        case "$fstype" in
            "vfat"|"fat32"|"fat16")
                # Partition 2: FAT32 EFI boot partition
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:11}"  # FAT32 has 11 character limit
                    if command -v fatlabel >/dev/null 2>&1; then
                        sudo fatlabel "$partition" "$new_label" 2>/dev/null || echo "Warning: Could not rename FAT partition $partition"
                    fi
                fi
                # Generate new UUID for FAT
                if command -v mlabel >/dev/null 2>&1; then
                    sudo mlabel -s -i "$partition" :: 2>/dev/null || echo "Warning: Could not change FAT serial on $partition"
                fi
                ;;
            "ext4")
                # Partition 3: ext4 root filesystem
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:16}"  # ext4 has 16 character limit
                    if command -v tune2fs >/dev/null 2>&1; then
                        sudo tune2fs -L "$new_label" "$partition" 2>/dev/null || echo "Warning: Could not rename ext4 partition $partition"
                    fi
                fi
                # Generate new UUID for ext4
                if command -v tune2fs >/dev/null 2>&1; then
                    sudo tune2fs -U random "$partition" 2>/dev/null || echo "Warning: Could not change UUID on $partition"
                fi
                ;;
            "")
                # Partition 1: No filesystem (skip)
                echo "Skipping partition $partition (no filesystem)"
                continue
                ;;
            *)
                echo "Warning: Unexpected filesystem type '$fstype' on $partition"
                continue
                ;;
        esac
        
        # Get the new UUID after regeneration
        local new_uuid
        new_uuid=$(lsblk -no UUID "$partition" 2>/dev/null)
        [[ -n "$new_uuid" ]] && new_uuids["$partition"]="$new_uuid"
        
        if [[ -n "$current_label" && -n "$new_label" ]]; then
            echo "Updated partition $partition: '$current_label' -> '$new_label'"
        fi
        if [[ -n "$current_uuid" && -n "$new_uuid" && "$current_uuid" != "$new_uuid" ]]; then
            echo "Regenerated UUID for $partition: $current_uuid -> $new_uuid"
        fi
    done
    
    # Return the UUID mappings for fstab update
    for partition in "${!old_uuids[@]}"; do
        if [[ -n "${new_uuids[$partition]}" ]]; then
            echo "UUID_MAPPING:${old_uuids[$partition]}:${new_uuids[$partition]}"
        fi
    done
}


# Welcome message
clear
echo "Stormux Gaming Image - USB to Disk Installer"
echo "This will clone the entire USB system to an internal disk."
echo

# Find source device
echo "Detecting source USB device..."
SOURCE_DEVICE=$(find_source_device)
if [[ -z "$SOURCE_DEVICE" ]]; then
    error_exit "Could not detect source USB device"
fi

SOURCE_SIZE=$(lsblk -dpno SIZE "$SOURCE_DEVICE" 2>/dev/null | tr -d ' ')
echo "Source device: $SOURCE_DEVICE ($SOURCE_SIZE)"
echo

# Detect target disks
echo "Detecting target disks..."
mapfile -t target_disks < <(detect_target_disks "$SOURCE_DEVICE")

if [[ ${#target_disks[@]} -eq 0 ]]; then
    error_exit "No suitable target disks found (excluding source USB)"
fi

# Display target disks
echo "Available target disks:"
for i in "${!target_disks[@]}"; do
    disk="${target_disks[$i]}"
    info=$(get_disk_info "$disk")
    echo "$((i+1)). $disk - $info"
done

# Get disk selection
while true; do
    echo
    echo "Enter the number of the disk to install to:"
    read -r selection
    
    if [[ "$selection" =~ ^[0-9]+$ ]] && [[ "$selection" -ge 1 ]] && [[ "$selection" -le ${#target_disks[@]} ]]; then
        TARGET_DEVICE="${target_disks[$((selection-1))]}"
        break
    else
        echo "Invalid selection. Please enter a number between 1 and ${#target_disks[@]}."
    fi
done

# Check target disk size
TARGET_SIZE_BYTES=$(lsblk -dpno SIZE -b "$TARGET_DEVICE" 2>/dev/null)
SOURCE_SIZE_BYTES=$(lsblk -dpno SIZE -b "$SOURCE_DEVICE" 2>/dev/null)

if [[ "$TARGET_SIZE_BYTES" -lt "$SOURCE_SIZE_BYTES" ]]; then
    error_exit "Target disk is smaller than source USB. Cannot proceed."
fi

# Final confirmation
target_info=$(get_disk_info "$TARGET_DEVICE")
echo
echo "FINAL WARNING:"
echo "Source: $SOURCE_DEVICE ($SOURCE_SIZE)"
echo "Target: $TARGET_DEVICE ($target_info)"
echo "ALL DATA ON THE TARGET DISK WILL BE PERMANENTLY DESTROYED!"
echo
echo "Type 'yes' to continue or any other key to cancel:"
read -r CONFIRM

if [[ "$CONFIRM" != "yes" ]]; then
    echo "Installation cancelled."
    exit 0
fi

# Rely on progress bar beeps for status
echo
echo "Fenrir will be silent during the install process except for progress beeps."
echo "To restore speech if needed, press any key."
echo "Press Enter to begin installation..."
read -r
echo "command tempdisablespeech" | socat - UNIX-CLIENT:/tmp/fenrirscreenreader-deamon.sock 2>/dev/null || true

# Unmount any mounted partitions on target disk
echo "Unmounting target disk partitions..."
sudo umount "${TARGET_DEVICE}"* 2>/dev/null || true

# Clone the USB to the target disk
echo "Cloning USB system to target disk..."
echo "This will take several minutes depending on USB size and disk speed."
echo

# Use dd to clone the entire device
if ! sudo dd if="$SOURCE_DEVICE" of="$TARGET_DEVICE" bs=4M oflag=sync status=progress; then
    error_exit "Failed to clone USB to target disk"
fi

# Sync to ensure all data is written
echo "Syncing data to disk..."
sudo sync

# Update the cloned system to remove USB-specific configurations
echo "Finalizing installation..."

# Mount the new root partition to make final adjustments
if [[ "$TARGET_DEVICE" =~ (nvme|mmcblk) ]]; then
    TARGET_ROOT_PART="${TARGET_DEVICE}p1"
else
    TARGET_ROOT_PART="${TARGET_DEVICE}1"
fi

# Find the actual root partition (might not be partition 1)
echo "Searching for ext4 root partition..."
for part in "${TARGET_DEVICE}"*; do
    if [[ "$part" != "$TARGET_DEVICE" ]]; then
        fstype=$(lsblk -no FSTYPE "$part" 2>/dev/null)
        echo "Checking $part: filesystem type = $fstype"
        if [[ "$fstype" == "ext4" ]]; then
            TARGET_ROOT_PART="$part"
            echo "Found root partition: $TARGET_ROOT_PART"
            break
        fi
    fi
done

# Regenerate partition UUIDs and labels to prevent boot conflicts with USB
echo "Regenerating partition identifiers..."
mapfile -t uuid_mappings < <(regenerate_partition_identifiers "$TARGET_DEVICE")
echo "Found ${#uuid_mappings[@]} UUID mappings to update"

# Mount and make final adjustments
TEMP_MOUNT="/mnt/stormux_target"
sudo mkdir -p "$TEMP_MOUNT"

echo "Mounting $TARGET_ROOT_PART for final adjustments..."
if sudo mount "$TARGET_ROOT_PART" "$TEMP_MOUNT" 2>/dev/null; then
    echo "Successfully mounted target partition"
    # Remove any USB-specific markers
    sudo rm -f "$TEMP_MOUNT/home/stormux/.firstboot" 2>/dev/null || true
    
    # Create .baremetal marker file to indicate installed system
    echo "Creating baremetal system marker..."
    sudo touch "$TEMP_MOUNT/home/stormux/.baremetal"
    sudo chown stormux:stormux "$TEMP_MOUNT/home/stormux/.baremetal" 2>/dev/null || true
    # Set immutable attribute to prevent accidental deletion
    sudo chattr +i "$TEMP_MOUNT/home/stormux/.baremetal" 2>/dev/null || echo "Warning: Could not set immutable attribute on .baremetal"
    
    # Update fstab with new UUIDs if any were generated
    if [[ ${#uuid_mappings[@]} -gt 0 ]]; then
        echo "Updating fstab with new UUIDs..."
        # Backup original fstab
        sudo cp "$TEMP_MOUNT/etc/fstab" "$TEMP_MOUNT/etc/fstab.backup" || echo "Warning: Could not backup fstab"
        
        # Apply UUID updates to fstab
        temp_fstab=$(mktemp)
        sudo cp "$TEMP_MOUNT/etc/fstab" "$temp_fstab"
        
        for mapping in "${uuid_mappings[@]}"; do
            if [[ "$mapping" =~ ^UUID_MAPPING:([^:]+):([^:]+)$ ]]; then
                old_uuid="${BASH_REMATCH[1]}"
                new_uuid="${BASH_REMATCH[2]}"
                
                # Replace old UUID with new UUID in fstab
                sed -i "s/UUID=$old_uuid/UUID=$new_uuid/g" "$temp_fstab"
                echo "Updated fstab: $old_uuid -> $new_uuid"
            fi
        done
        
        # Copy updated fstab back
        sudo cp "$temp_fstab" "$TEMP_MOUNT/etc/fstab"
        rm -f "$temp_fstab"
        echo "fstab updated successfully"
    else
        echo "No UUID mappings found - keeping original fstab"
        sudo cp "$TEMP_MOUNT/etc/fstab" "$TEMP_MOUNT/etc/fstab.backup" || echo "Warning: Could not backup fstab"
    fi
    
    # Also mount and update GRUB config on EFI partition
    EFI_MOUNT="/mnt/stormux_efi"
    sudo mkdir -p "$EFI_MOUNT"
    
    # Find EFI partition (should be partition 2)
    if [[ "$TARGET_DEVICE" =~ (nvme|mmcblk) ]]; then
        EFI_PART="${TARGET_DEVICE}p2"
    else
        EFI_PART="${TARGET_DEVICE}2"
    fi
    
    if sudo mount "$EFI_PART" "$EFI_MOUNT" 2>/dev/null; then
        echo "Updating GRUB configuration with new UUIDs..."
        if [[ -f "$EFI_MOUNT/grub/grub.cfg" && ${#uuid_mappings[@]} -gt 0 ]]; then
            # Backup original grub config
            sudo cp "$EFI_MOUNT/grub/grub.cfg" "$EFI_MOUNT/grub/grub.cfg.backup" || echo "Warning: Could not backup grub.cfg"
            
            # Apply UUID updates to grub config
            grub_temp=$(mktemp)
            sudo cp "$EFI_MOUNT/grub/grub.cfg" "$grub_temp"
            
            for mapping in "${uuid_mappings[@]}"; do
                if [[ "$mapping" =~ ^UUID_MAPPING:([^:]+):([^:]+)$ ]]; then
                    grub_old_uuid="${BASH_REMATCH[1]}"
                    grub_new_uuid="${BASH_REMATCH[2]}"
                    
                    # Replace old UUID with new UUID in grub config
                    sed -i "s/${grub_old_uuid}/${grub_new_uuid}/g" "$grub_temp"
                    echo "Updated grub.cfg: $grub_old_uuid -> $grub_new_uuid"
                fi
            done
            
            # Copy updated grub config back
            sudo cp "$grub_temp" "$EFI_MOUNT/grub/grub.cfg"
            rm -f "$grub_temp"
            echo "GRUB configuration updated successfully"
        else
            echo "No GRUB config found or no UUID mappings to update"
        fi
        
        sudo umount "$EFI_MOUNT"
    else
        echo "Warning: Could not mount EFI partition $EFI_PART"
    fi
    sudo rmdir "$EFI_MOUNT"

# Restore speech
echo "command toggletempdisablespeech" | socat - UNIX-CLIENT:/tmp/fenrirscreenreader-deamon.sock

    echo "Unmounting target partition..."
    sudo umount "$TEMP_MOUNT"
else
    echo "ERROR: Failed to mount target partition $TARGET_ROOT_PART"
    echo "Installation may have failed - please check the target disk"
    echo "However, the disk cloning operation completed."
fi
sudo rmdir "$TEMP_MOUNT"

# Success message (using Fenrir auto-restore pattern)
echo
echo "========================================="
echo "Installation completed successfully!"
echo "The USB system has been cloned to $TARGET_DEVICE"
echo "You can now reboot and remove the USB drive."
echo "The system will boot from the internal disk."
echo "========================================="
echo
echo "Press enter to continue..."
read -r