#!/usr/bin/env bash

# Stormux Gaming Image - Rsync-based Disk Installer
# Copies the live system to a target disk using rsync

set -e

# Text output functions
speak_and_echo() {
    echo "$1"
    spd-say "$1" 2>/dev/null || true
}

error_exit() {
    speak_and_echo "Error: $1"
    exit 1
}

# 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"
}

# Welcome message
clear
speak_and_echo "Stormux Gaming Image Disk Installer"
speak_and_echo "This will install the current live system to a hard disk."
echo
speak_and_echo "WARNING: This will completely erase the target disk!"

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

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

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

# Get disk selection
while true; do
    echo
    speak_and_echo "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
        speak_and_echo "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")

echo
speak_and_echo "FINAL WARNING:"
speak_and_echo "You have selected: $selected_disk ($selected_size, $selected_model)"
speak_and_echo "ALL DATA ON THIS DISK WILL BE PERMANENTLY DESTROYED!"
echo
speak_and_echo "Type 'yes' to continue or any other key to cancel:"
read -r CONFIRM

if [[ "$CONFIRM" != "yes" ]]; then
    speak_and_echo "Installation cancelled."
    exit 0
fi

# Check if we're on UEFI or BIOS
if [[ -d /sys/firmware/efi ]]; then
    BOOT_MODE="UEFI"
    speak_and_echo "Detected UEFI boot mode"
else
    BOOT_MODE="BIOS"
    speak_and_echo "Detected BIOS boot mode"
fi

speak_and_echo "Starting installation to $selected_disk..."

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

# Create partition table
speak_and_echo "Creating partition table..."
if [[ "$BOOT_MODE" == "UEFI" ]]; then
    # UEFI: GPT with EFI system partition
    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%
    
    # Set partition variables
    if [[ "$selected_disk" =~ nvme|mmcblk ]]; then
        EFI_PART="${selected_disk}p1"
        ROOT_PART="${selected_disk}p2"
    else
        EFI_PART="${selected_disk}1"
        ROOT_PART="${selected_disk}2"
    fi
else
    # BIOS: MBR with boot flag
    parted -s "$selected_disk" mklabel msdos
    parted -s "$selected_disk" mkpart primary ext4 1MiB 100%
    parted -s "$selected_disk" set 1 boot on
    
    # Set partition variables
    if [[ "$selected_disk" =~ nvme|mmcblk ]]; then
        ROOT_PART="${selected_disk}p1"
    else
        ROOT_PART="${selected_disk}1"
    fi
fi

# Wait for kernel to recognize partitions
sleep 2
partprobe "$selected_disk"
sleep 2

# Format partitions
speak_and_echo "Formatting partitions..."
if [[ "$BOOT_MODE" == "UEFI" ]]; then
    mkfs.fat -F32 -n "STORMUX-EFI" "$EFI_PART"
fi
mkfs.ext4 -F -L "STORMUX-ROOT" "$ROOT_PART"

# Create mount points
MOUNT_ROOT="/mnt/install_target"
mkdir -p "$MOUNT_ROOT"

# Mount root partition
speak_and_echo "Mounting partitions..."
mount "$ROOT_PART" "$MOUNT_ROOT"

# Mount EFI partition if UEFI
if [[ "$BOOT_MODE" == "UEFI" ]]; then
    mkdir -p "$MOUNT_ROOT/boot/efi"
    mount "$EFI_PART" "$MOUNT_ROOT/boot/efi"
fi

# Copy system using rsync
speak_and_echo "Copying system files... This may take several minutes."
rsync -avx \
    --exclude=/tmp/* \
    --exclude=/proc/* \
    --exclude=/sys/* \
    --exclude=/dev/* \
    --exclude=/run/* \
    --exclude=/var/tmp/* \
    --exclude=/var/log/journal/* \
    --exclude=/mnt/* \
    --exclude=/media/* \
    --exclude=/lost+found \
    --exclude=/home/stormux/.cache/* \
    --exclude=/home/stormux/Downloads/* \
    / "$MOUNT_ROOT/"

# Create necessary directories
speak_and_echo "Creating system directories..."
mkdir -p "$MOUNT_ROOT"/{tmp,proc,sys,dev,run,var/tmp,var/log/journal,mnt,media}

# Generate fstab
speak_and_echo "Generating fstab..."
ROOT_UUID=$(blkid -s UUID -o value "$ROOT_PART")
cat > "$MOUNT_ROOT/etc/fstab" << EOF
# Static information about the filesystems.
# See fstab(5) for details.

# <file system> <dir> <type> <options> <dump> <pass>
UUID=$ROOT_UUID / ext4 rw,relatime 0 1
EOF

if [[ "$BOOT_MODE" == "UEFI" ]]; then
    EFI_UUID=$(blkid -s UUID -o value "$EFI_PART")
    echo "UUID=$EFI_UUID /boot/efi vfat rw,relatime,fmask=0022,dmask=0022,codepage=437,iocharset=ascii,shortname=mixed,utf8,errors=remount-ro 0 2" >> "$MOUNT_ROOT/etc/fstab"
fi

# Install bootloader
speak_and_echo "Installing bootloader..."
arch-chroot "$MOUNT_ROOT" /bin/bash << EOF
# Update package database and install GRUB
pacman -Sy --noconfirm

if [[ "$BOOT_MODE" == "UEFI" ]]; then
    # UEFI installation
    pacman -S --needed --noconfirm grub efibootmgr
    grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=Stormux
else
    # BIOS installation  
    pacman -S --needed --noconfirm grub
    grub-install --target=i386-pc "$selected_disk"
fi

# Generate GRUB configuration with accessibility
grub-mkconfig -o /boot/grub/grub.cfg

# Enable accessibility in GRUB by default
if ! grep -q "GRUB_CMDLINE_LINUX_DEFAULT.*accessibility=on" /etc/default/grub; then
    sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="/GRUB_CMDLINE_LINUX_DEFAULT="accessibility=on /' /etc/default/grub
    grub-mkconfig -o /boot/grub/grub.cfg
fi

# Enable essential services
systemctl enable NetworkManager
systemctl enable pipewire.service
systemctl enable pipewire-pulse.service
systemctl enable unmute-audio.service
systemctl enable shutdown-sound.service
systemctl enable fenrirscreenreader-tty.service
systemctl disable espeakup.service

# Set up user services
systemctl --global enable pipewire.service
systemctl --global enable pipewire-pulse.service
EOF

# Remove firstboot marker if it exists
rm -f "$MOUNT_ROOT/home/stormux/.firstboot"

# Clean up mount points
speak_and_echo "Cleaning up..."
if [[ "$BOOT_MODE" == "UEFI" ]]; then
    umount "$MOUNT_ROOT/boot/efi"
fi
umount "$MOUNT_ROOT"
rmdir "$MOUNT_ROOT"

# Success message
speak_and_echo "Installation completed successfully!"
echo
speak_and_echo "The system has been installed to $selected_disk"
speak_and_echo "You can now reboot and remove the USB drive."
echo
speak_and_echo "Press Enter to continue..."
read -r