#!/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
        # 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=$(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"
}

# Function to rename partition labels to prevent boot conflicts
rename_partition_labels() {
    local target_device="$1"
    echo "Renaming partition labels to prevent boot conflicts..."
    
    # 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}")
    
    # Rename each partition based on filesystem type
    for partition in "${partitions[@]}"; do
        local fstype
        local current_label
        fstype=$(lsblk -no FSTYPE "$partition" 2>/dev/null)
        current_label=$(lsblk -no LABEL "$partition" 2>/dev/null)
        
        # Skip if no filesystem or no current label
        [[ -z "$fstype" || -z "$current_label" ]] && continue
        
        # Determine new label based on current label and add -HDD suffix
        local new_label=""
        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
        
        # Truncate label if too long (filesystem limits)
        case "$fstype" in
            "vfat"|"fat32"|"fat16")
                # FAT32 has 11 character limit
                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
                ;;
            "ext2"|"ext3"|"ext4")
                # ext* has 16 character limit
                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
                ;;
            "ntfs")
                # NTFS has 32 character limit
                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
                ;;
            "xfs")
                # XFS has 12 character limit
                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
                ;;
            *)
                echo "Info: Skipping unknown filesystem type '$fstype' on $partition"
                ;;
        esac
        
        echo "Renamed partition $partition: '$current_label' -> '$new_label'"
    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
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

# 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
    
    sudo umount "$TEMP_MOUNT"
fi
sudo rmdir "$TEMP_MOUNT"

# Rename all partition labels to prevent boot conflicts with USB
rename_partition_labels "$TARGET_DEVICE"

# 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