#!/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

# 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
        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
    
    # Remove STORMUX label from filesystem to prevent false USB detection
    if command -v tune2fs >/dev/null 2>&1; then
        sudo tune2fs -L "" "$TARGET_ROOT_PART" 2>/dev/null || true
    fi
    
    # Update any USB-specific configurations if needed
    
    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