#!/bin/bash

# Stormux Gaming Image Disk Installer
# WARNING: This will completely wipe the selected disk!

set -e

# Function to print messages
print_message() {
    echo "$1"
}

# Function to get user confirmation
get_confirmation() {
    local prompt="$1"
    local required_response="$2"
    local response
    
    while true; do
        print_message "$prompt"
        read -r response
        if [[ "$response" == "$required_response" ]]; then
            return 0
        else
            print_message "Response must be exactly: $required_response"
        fi
    done
}

# Function to detect available disks
detect_disks() {
    local disks=()
    
    # Find all block devices that are disks (not partitions)
    while IFS= read -r disk; do
        # Skip if it's a partition, loop device, or CD-ROM
        if [[ ! "$disk" =~ [0-9]$ ]] && [[ ! "$disk" =~ ^/dev/loop ]] && [[ ! "$disk" =~ ^/dev/sr ]]; then
            # Check if it's actually a disk
            if [[ -b "$disk" ]]; then
                disks+=("$disk")
            fi
        fi
    done < <(lsblk -dpno NAME 2>/dev/null)
    
    printf '%s\n' "${disks[@]}"
}

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

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

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    print_message "This script must be run as root."
    print_message "Please run: sudo $0"
    exit 1
fi

# Welcome message
print_message ""
print_message "=== Stormux Gaming Image Installer ==="
print_message "WARNING: This installer will COMPLETELY WIPE the selected disk!"
print_message "All data on the target disk will be permanently lost."

# Detect available disks
print_message ""
print_message "Detecting available disks..."
mapfile -t available_disks < <(detect_disks)

if [[ ${#available_disks[@]} -eq 0 ]]; then
    print_message "No suitable disks found!"
    exit 1
fi

# Display available disks
print_message ""
print_message "Available disks:"
for i in "${!available_disks[@]}"; do
    disk="${available_disks[$i]}"
    size=$(get_disk_size "$disk")
    model=$(get_disk_model "$disk")
    print_message "$((i+1)). $disk - $size - $model" "$((i+1)). $disk - $size - $model"
done

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

# Final confirmation
selected_size=$(get_disk_size "$selected_disk")
selected_model=$(get_disk_model "$selected_disk")

print_message ""
print_message "FINAL WARNING:"
print_message "You have selected: $selected_disk ($selected_size, $selected_model)"
print_message "ALL DATA ON THIS DISK WILL BE PERMANENTLY DESTROYED!"

get_confirmation "Type 'DESTROY ALL DATA' to continue (case sensitive): " "DESTROY ALL DATA"

print_message ""
print_message "Starting installation to $selected_disk..."

# Unmount any existing partitions on the target disk
print_message "Unmounting any existing partitions..."
umount "${selected_disk}"* 2>/dev/null || true

# Create partition table and partitions
print_message "Creating partition table..."
parted -s "$selected_disk" mklabel gpt
parted -s "$selected_disk" mkpart primary fat32 1MiB 513MiB
parted -s "$selected_disk" set 1 esp on
parted -s "$selected_disk" mkpart primary ext4 513MiB 100%

# Get partition names
if [[ "$selected_disk" =~ nvme|mmcblk ]]; then
    boot_partition="${selected_disk}p1"
    root_partition="${selected_disk}p2"
else
    boot_partition="${selected_disk}1"
    root_partition="${selected_disk}2"
fi

# Format partitions
print_message "Formatting partitions..."
mkfs.fat -F32 "$boot_partition"
mkfs.ext4 -F "$root_partition"

# Create mount points and mount
mkdir -p /mnt/stormux-install
mount "$root_partition" /mnt/stormux-install
mkdir -p /mnt/stormux-install/boot
mount "$boot_partition" /mnt/stormux-install/boot

# Copy system files
print_message "Copying system files... This may take several minutes."
if [[ -f /run/archiso/img_dev ]]; then
    # Extract from squashfs
    unsquashfs -f -d /mnt/stormux-install /run/archiso/sfs/airootfs/airootfs.sfs
else
    # Fallback: copy from live system
    rsync -aHAXS --exclude={"/dev/*","/proc/*","/sys/*","/tmp/*","/run/*","/mnt/*","/media/*","/lost+found"} / /mnt/stormux-install/
fi

# Configure for installed system
print_message "Configuring system for installed boot..."

# Remove live-specific files and services
rm -f /mnt/stormux-install/etc/systemd/system/getty@tty1.service.d/autologin.conf
rm -f /mnt/stormux-install/etc/systemd/system/multi-user.target.wants/choose-mirror.service
rm -rf /mnt/stormux-install/etc/systemd/system/etc-pacman.d-gnupg.mount

# Generate fstab
genfstab -U /mnt/stormux-install >> /mnt/stormux-install/etc/fstab

# Install and configure GRUB
print_message "Installing bootloader..."
arch-chroot /mnt/stormux-install grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=stormux
arch-chroot /mnt/stormux-install grub-mkconfig -o /boot/grub/grub.cfg

# Set hostname
echo "stormux-gaming" > /mnt/stormux-install/etc/hostname

# Enable necessary services
arch-chroot /mnt/stormux-install systemctl enable NetworkManager
arch-chroot /mnt/stormux-install systemctl enable pipewire pipewire-pulse
arch-chroot /mnt/stormux-install systemctl enable speech-dispatcher
arch-chroot /mnt/stormux-install systemctl enable fenrirscreenreader

# Cleanup
umount -R /mnt/stormux-install
rmdir /mnt/stormux-install

print_message ""
print_message "Installation completed successfully!"
print_message "You can now reboot and remove the installation media."
print_message "The system will boot directly into the Stormux Gaming environment."