#!/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=$(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=$(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
        if [[ ! "$disk" =~ [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=$(lsblk -dpno SIZE "$disk" 2>/dev/null | tr -d ' ')
    local model=$(lsblk -dpno MODEL "$disk" 2>/dev/null | tr -d ' ' || echo "Unknown")
    echo "$size - $model"
}

# 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

# 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 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
        local fstype=$(lsblk -no FSTYPE "$part" 2>/dev/null)
        if [[ "$fstype" == "ext4" ]]; then
            TARGET_ROOT_PART="$part"
            break
        fi
    fi
done

# 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
    
    # Update any USB-specific configurations if needed
    # (Add any specific cleanups here)
    
    sudo umount "$TEMP_MOUNT"
fi
sudo rmdir "$TEMP_MOUNT"

# 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