#!/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" ]]; then
            partitions+=("$partition")
        fi
    done < <(lsblk -lpno NAME "$target_device" 2>/dev/null | grep "^${target_device}")
    
    # 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
        
        # Generate new UUID and update label based on filesystem type
        case "$fstype" in
            "vfat"|"fat32"|"fat16")
                # FAT32 has 11 character limit for labels
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:11}"
                    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 - mlabel -s generates a new random serial
                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
                ;;
            "ext2"|"ext3"|"ext4")
                # ext* has 16 character limit for labels
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:16}"
                    if command -v tune2fs >/dev/null 2>&1; then
                        sudo tune2fs -L "$new_label" "$partition" 2>/dev/null || echo "Warning: Could not rename ext partition $partition"
                    fi
                fi
                # Generate new UUID for ext filesystems
                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
                ;;
            "ntfs")
                # NTFS has 32 character limit for labels
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:32}"
                    if command -v ntfslabel >/dev/null 2>&1; then
                        sudo ntfslabel "$partition" "$new_label" 2>/dev/null || echo "Warning: Could not rename NTFS partition $partition"
                    fi
                fi
                ;;
            "xfs")
                # XFS has 12 character limit for labels
                if [[ -n "$new_label" ]]; then
                    new_label="${new_label:0:12}"
                    if command -v xfs_admin >/dev/null 2>&1; then
                        sudo xfs_admin -L "$new_label" "$partition" 2>/dev/null || echo "Warning: Could not rename XFS partition $partition"
                    fi
                fi
                # Generate new UUID for XFS
                if command -v xfs_admin >/dev/null 2>&1; then
                    sudo xfs_admin -U generate "$partition" 2>/dev/null || echo "Warning: Could not change UUID on $partition"
                fi
                ;;
            *)
                echo "Info: Skipping unknown 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
}

# Function to update fstab with new UUIDs
update_fstab_uuids() {
    local mount_point="$1"
    shift
    local uuid_mappings=("$@")
    
    local fstab_file="$mount_point/etc/fstab"
    
    if [[ ! -f "$fstab_file" ]]; then
        echo "Warning: /etc/fstab not found in mounted system"
        return 1
    fi
    
    echo "Updating /etc/fstab with new UUIDs..."
    
    # Create backup of original fstab
    sudo cp "$fstab_file" "$fstab_file.backup" || {
        echo "Warning: Could not create fstab backup"
        return 1
    }
    
    # Apply UUID mappings to fstab
    local temp_fstab
    temp_fstab=$(mktemp)
    
    for mapping in "${uuid_mappings[@]}"; do
        if [[ "$mapping" =~ ^UUID_MAPPING:([^:]+):([^:]+)$ ]]; then
            local old_uuid="${BASH_REMATCH[1]}"
            local new_uuid="${BASH_REMATCH[2]}"
            
            # Replace old UUID with new UUID in fstab
            sudo sed "s/UUID=$old_uuid/UUID=$new_uuid/g" "$fstab_file" | sudo tee "$temp_fstab" > /dev/null
            sudo cp "$temp_fstab" "$fstab_file"
            
            echo "Updated fstab: $old_uuid -> $new_uuid"
        fi
    done
    
    rm -f "$temp_fstab"
    return 0
}

# 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
TARGET_ROOT_PART="${TARGET_DEVICE}1"
if [[ "$TARGET_DEVICE" =~ nvme|mmcblk ]]; then
    TARGET_ROOT_PART="${TARGET_DEVICE}p1"
fi

# Find the actual root partition (might not be partition 1)
for part in "${TARGET_DEVICE}"*; do
    if [[ "$part" != "$TARGET_DEVICE" ]]; then
        fstype=$(lsblk -no FSTYPE "$part" 2>/dev/null)
        if [[ "$fstype" == "ext4" ]]; then
            TARGET_ROOT_PART="$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")

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

if sudo mount "$TARGET_ROOT_PART" "$TEMP_MOUNT" 2>/dev/null; then
    # 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 /etc/fstab with new UUIDs
    if [[ ${#uuid_mappings[@]} -gt 0 ]]; then
        update_fstab_uuids "$TEMP_MOUNT" "${uuid_mappings[@]}"
    fi
    
    sudo umount "$TEMP_MOUNT"
fi
sudo rmdir "$TEMP_MOUNT"

# Re-enable speech before success message
echo "command toggletempdisablespeech" | socat - UNIX-CLIENT:/tmp/fenrirscreenreader-deamon.sock 2>/dev/null || true

# Success message
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 "Press Enter to continue..."
read -r