#!/usr/bin/env bash

# Stormux Gaming Image - Ultra-Reliable Disk Installer
# Clones USB system to internal disk with proper bootloader installation

set -euo pipefail

# Initialize logging
LOGDIR="/home/stormux/Logs"
LOGFILE="$LOGDIR/install_to_disk.log"
mkdir -p "$LOGDIR"

# Logging function
log() {
    echo "$*" | tee -a "$LOGFILE"
}

log_error() {
    echo "[ERROR] $*" | tee -a "$LOGFILE" >&2
}

cleanup_mounts() {
    # Clean up any leftover mounts
    local mount_point="${1:-/mnt/stormux_target}"
    if [[ -d "$mount_point" ]]; then
        sudo umount "$mount_point/boot" 2>/dev/null || true
        sudo umount "$mount_point/proc" 2>/dev/null || true
        sudo umount "$mount_point/sys" 2>/dev/null || true
        sudo umount "$mount_point/dev" 2>/dev/null || true
        sudo umount "$mount_point" 2>/dev/null || true
        sudo rmdir "$mount_point" 2>/dev/null || true
    fi
}

error_exit() {
    log_error "$1"
    echo "Installation failed. Log file: $LOGFILE"
    cleanup_mounts
    restore_speech
    echo
    read -rp "Press enter to continue..."
    exit 1
}

# Speech management
disable_speech() {
    echo "command tempdisablespeech" | socat - UNIX-CLIENT:/tmp/fenrirscreenreader-deamon.sock 2>/dev/null || true
}

restore_speech() {
    echo "command toggletempdisablespeech" | socat - UNIX-CLIENT:/tmp/fenrirscreenreader-deamon.sock 2>/dev/null || true
}

# Function to find the source USB device
find_source_device() {
    local root_device
    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
        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 partitions, loop devices, CD-ROMs, and source device
        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
    local model
    size=$(lsblk -dpno SIZE "$disk" 2>/dev/null | tr -d ' ')
    model=$(lsblk -dpno MODEL "$disk" 2>/dev/null | tr -d ' ' || echo "Unknown")
    echo "$size - $model"
}

# Function to detect partitions by filesystem type
# shellcheck disable=SC2154
detect_partitions() {
    local device="$1"
    # shellcheck disable=SC2178
    local -n result=$2  # nameref to associative array

    log "Detecting partition structure on $device..."

    # Get all partitions
    while IFS= read -r line; do
        local part fstype label
        part=$(echo "$line" | awk '{print $1}')
        # Remove lsblk tree characters that break mount commands
        part="${part//[├─└│]/}"
        # Add /dev/ prefix if not present
        [[ "$part" != /dev/* ]] && part="/dev/$part"
        fstype=$(echo "$line" | awk '{print $2}')
        label=$(echo "$line" | awk '{print $3}')

        # Skip if it's the device itself, not a partition
        [[ "$part" == "$device" ]] && continue

        case "$fstype" in
            "")
                # BIOS boot partition (no filesystem)
                result[bios]="$part"
                log "  BIOS boot partition: $part"
                ;;
            "vfat")
                # EFI partition
                result[efi]="$part"
                log "  EFI partition: $part (label: ${label:-none})"
                ;;
            "ext4")
                # Root partition
                result[root]="$part"
                log "  Root partition: $part (label: ${label:-none})"
                ;;
            *)
                log "  Unknown partition type: $part ($fstype)"
                ;;
        esac
    done < <(lsblk -no NAME,FSTYPE,LABEL "$device" 2>/dev/null | tail -n +2)

    # Validate we found all required partitions
    if [[ -z "${result[root]:-}" ]]; then
        return 1
    fi

    return 0
}

# Function to regenerate UUIDs and update labels
regenerate_partition_identifiers() {
    local device="$1"
    local -n partitions=$2
    declare -A uuid_map

    log "Regenerating partition UUIDs and labels..."

    # Process EFI partition
    if [[ -n "${partitions[efi]:-}" ]]; then
        local efi_part="${partitions[efi]}"
        local old_uuid
        old_uuid=$(lsblk -no UUID "$efi_part" 2>/dev/null || echo "")

        # Change label
        if command -v fatlabel >/dev/null 2>&1; then
            sudo fatlabel "$efi_part" "BOOT-HDD" 2>/dev/null || log "Warning: Could not rename EFI partition"
        fi

        # Generate new UUID for FAT
        if command -v mlabel >/dev/null 2>&1; then
            sudo mlabel -s -i "$efi_part" :: 2>/dev/null || log "Warning: Could not change FAT serial"
        fi

        local new_uuid
        new_uuid=$(lsblk -no UUID "$efi_part" 2>/dev/null || echo "")

        if [[ -n "$old_uuid" && -n "$new_uuid" ]]; then
            uuid_map["$old_uuid"]="$new_uuid"
            log "  EFI UUID: $old_uuid -> $new_uuid"
        fi
    fi

    # Process root partition
    if [[ -n "${partitions[root]:-}" ]]; then
        local root_part="${partitions[root]}"
        local old_uuid
        old_uuid=$(lsblk -no UUID "$root_part" 2>/dev/null || echo "")

        # Change label and UUID
        if command -v tune2fs >/dev/null 2>&1; then
            sudo tune2fs -L "STORMUX-HDD" "$root_part" 2>/dev/null || log "Warning: Could not rename root partition"
            sudo tune2fs -U random "$root_part" 2>/dev/null || log "Warning: Could not change root UUID"
        fi

        local new_uuid
        new_uuid=$(lsblk -no UUID "$root_part" 2>/dev/null || echo "")

        if [[ -n "$old_uuid" && -n "$new_uuid" ]]; then
            uuid_map["$old_uuid"]="$new_uuid"
            log "  Root UUID: $old_uuid -> $new_uuid"
        fi
    fi

    # Return UUID mappings
    for old in "${!uuid_map[@]}"; do
        echo "UUID_MAP:$old:${uuid_map[$old]}"
    done
}

# Function to update fstab with new UUIDs
update_fstab() {
    local mount_point="$1"
    shift
    local mappings=("$@")

    log "Updating fstab with new UUIDs..."

    local fstab="$mount_point/etc/fstab"

    if [[ ! -f "$fstab" ]]; then
        log_error "fstab not found at $fstab"
        return 1
    fi

    # Backup original
    sudo cp "$fstab" "$fstab.backup" || log "Warning: Could not backup fstab"

    # Apply UUID mappings
    local temp_fstab
    temp_fstab=$(mktemp)
    sudo cp "$fstab" "$temp_fstab"

    for mapping in "${mappings[@]}"; do
        if [[ "$mapping" =~ ^UUID_MAP:([^:]+):([^:]+)$ ]]; then
            local old_uuid="${BASH_REMATCH[1]}"
            local new_uuid="${BASH_REMATCH[2]}"

            sed -i "s/UUID=$old_uuid/UUID=$new_uuid/g" "$temp_fstab"
            log "  Updated fstab: $old_uuid -> $new_uuid"
        fi
    done

    sudo cp "$temp_fstab" "$fstab"
    rm -f "$temp_fstab"

    log "fstab updated successfully"
    return 0
}

# Function to install GRUB in chroot
install_grub() {
    local mount_point="$1"
    local target_device="$2"
    local -n parts=$3

    log "Installing GRUB bootloader to $target_device..."

    # Ensure EFI partition is mounted inside chroot
    if [[ -n "${parts[efi]:-}" ]]; then
        sudo mkdir -p "$mount_point/boot"
        if ! mountpoint -q "$mount_point/boot"; then
            sudo mount "${parts[efi]}" "$mount_point/boot" || error_exit "Failed to mount EFI partition"
            log "  Mounted EFI partition at /boot"
        fi
    fi

    # Bind mount necessary filesystems for chroot
    sudo mount --bind /dev "$mount_point/dev" || error_exit "Failed to bind mount /dev"
    sudo mount --bind /sys "$mount_point/sys" || error_exit "Failed to bind mount /sys"
    sudo mount --bind /proc "$mount_point/proc" || error_exit "Failed to bind mount /proc"
    log "  Prepared chroot environment"

    # Install GRUB for BIOS (allow to fail gracefully)
    log "  Installing GRUB for BIOS boot (warnings expected on UEFI systems)..."
    if sudo arch-chroot "$mount_point" grub-install --target=i386-pc --recheck "$target_device" 2>&1 | tee -a "$LOGFILE"; then
        log "  BIOS boot installation succeeded"
    else
        log "  BIOS boot installation completed with warnings (expected on UEFI systems)"
    fi

    # Install GRUB for UEFI (must succeed)
    # Use --removable flag to install to default EFI fallback location
    log "  Installing GRUB for UEFI boot..."
    if ! sudo arch-chroot "$mount_point" grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=STORMUX-HDD --recheck --removable 2>&1 | tee -a "$LOGFILE"; then
        error_exit "UEFI GRUB installation failed"
    fi
    log "  UEFI boot installation succeeded"

    # Verify /etc/default/grub exists before generating config
    if [[ ! -f "$mount_point/etc/default/grub" ]]; then
        log_error "/etc/default/grub not found on target system"
        error_exit "Cannot generate GRUB config without /etc/default/grub"
    fi

    # Generate GRUB configuration
    log "  Generating GRUB configuration..."
    if ! sudo arch-chroot "$mount_point" grub-mkconfig -o /boot/grub/grub.cfg 2>&1 | tee -a "$LOGFILE"; then
        error_exit "GRUB configuration generation failed"
    fi
    log "  GRUB configuration generated successfully"

    # Unmount bind mounts
    sudo umount "$mount_point/proc" 2>/dev/null || true
    sudo umount "$mount_point/sys" 2>/dev/null || true
    sudo umount "$mount_point/dev" 2>/dev/null || true

    log "GRUB installation completed successfully"
    return 0
}

# Function to validate installation
validate_installation() {
    local mount_point="$1"

    log "Validating installation..."

    # Check fstab exists and has valid entries
    if [[ ! -f "$mount_point/etc/fstab" ]]; then
        log_error "fstab not found"
        return 1
    fi

    # Check GRUB config exists
    if [[ ! -f "$mount_point/boot/grub/grub.cfg" ]]; then
        log_error "GRUB configuration not found"
        return 1
    fi

    # Check for baremetal marker
    if [[ ! -f "$mount_point/home/stormux/.baremetal" ]]; then
        log_error "Baremetal marker not found"
        return 1
    fi

    log "Validation passed"
    return 0
}

###################
# Main Installation
###################

clear
log "========================================"
log "Stormux Gaming Image - Disk Installer"
log "========================================"
log "Log file: $LOGFILE"
echo
echo "This will clone the USB system to an internal disk."
echo

# Find source device
log "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 ' ')
log "Source device: $SOURCE_DEVICE ($SOURCE_SIZE)"
echo "Source device: $SOURCE_DEVICE ($SOURCE_SIZE)"
echo

# Detect target disks
log "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"
fi

# Display target disks
echo "Available target disks:"
log "Available target disks:"
for i in "${!target_disks[@]}"; do
    disk="${target_disks[$i]}"
    info=$(get_disk_info "$disk")
    echo "$((i+1)). $disk - $info"
    log "  $((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

log "Selected target device: $TARGET_DEVICE"

# Safety check: ensure target != source
if [[ "$TARGET_DEVICE" == "$SOURCE_DEVICE" ]]; then
    error_exit "Target device cannot be the same as source device"
fi

# 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"
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
log "Final confirmation - Target: $TARGET_DEVICE"
echo "Type 'yes' to continue or any other key to cancel:"
read -r CONFIRM

if [[ "$CONFIRM" != "yes" ]]; then
    log "Installation cancelled by user"
    echo "Installation cancelled."
    exit 0
fi

log "User confirmed installation"

# Disable speech during installation
echo
echo "Fenrir will be silent during installation except for progress beeps."
echo "Press Enter to begin..."
read -r
disable_speech

# Unmount any mounted partitions on target disk
log "Unmounting target disk partitions..."
sudo umount "${TARGET_DEVICE}"* 2>/dev/null || true

# Clone the USB to target disk
log "Starting disk clone operation..."
echo "Cloning USB system to target disk..."
echo "This will take several minutes depending on USB size and disk speed."
echo

if ! sudo dd if="$SOURCE_DEVICE" of="$TARGET_DEVICE" bs=4M oflag=sync status=progress 2>&1 | tee -a "$LOGFILE"; then
    error_exit "Failed to clone USB to target disk"
fi

log "Disk clone completed successfully"

# Sync and refresh partition table
log "Syncing data to disk..."
sudo sync

log "Refreshing partition table..."
# Use -s flag for script mode, pipe 'Fix' response for GPT expansion
echo "Fix" | sudo partprobe -s "$TARGET_DEVICE" 2>&1 | tee -a "$LOGFILE" || log "Warning: partprobe reported issues"
sudo udevadm settle --timeout=10 || log "Warning: udevadm settle timeout"
sleep 2

# Detect partition structure
log "Analyzing cloned partition structure..."
declare -A target_partitions
if ! detect_partitions "$TARGET_DEVICE" target_partitions; then
    error_exit "Failed to detect partition structure on target disk"
fi

# Validate partition structure
if [[ -z "${target_partitions[root]:-}" ]]; then
    error_exit "Could not find root partition on target disk"
fi

log "Partition detection complete:"
log "  BIOS: ${target_partitions[bios]:-none}"
log "  EFI: ${target_partitions[efi]:-none}"
log "  Root: ${target_partitions[root]}"

# Regenerate UUIDs and labels BEFORE mounting
log "Regenerating partition identifiers..."
mapfile -t uuid_mappings < <(regenerate_partition_identifiers "$TARGET_DEVICE" target_partitions)
log "Generated ${#uuid_mappings[@]} UUID mappings"

# Trigger udev to update with new UUIDs
log "Updating system device database..."
sudo udevadm trigger --subsystem-match=block
sudo udevadm settle --timeout=10 || log "Warning: udevadm settle timeout"
sleep 1

# Mount root partition
TEMP_MOUNT="/mnt/stormux_target"
sudo mkdir -p "$TEMP_MOUNT"

log "Mounting root partition ${target_partitions[root]}..."
if ! sudo mount "${target_partitions[root]}" "$TEMP_MOUNT" 2>&1 | tee -a "$LOGFILE"; then
    error_exit "Failed to mount root partition"
fi

log "Root partition mounted successfully"

# Update fstab with new UUIDs
if ! update_fstab "$TEMP_MOUNT" "${uuid_mappings[@]}"; then
    error_exit "Failed to update fstab"
fi

# Create baremetal marker
log "Creating baremetal system marker..."
sudo touch "$TEMP_MOUNT/home/stormux/.baremetal"
sudo chown 1000:1000 "$TEMP_MOUNT/home/stormux/.baremetal" 2>/dev/null || true
sudo chattr +i "$TEMP_MOUNT/home/stormux/.baremetal" 2>/dev/null || log "Warning: Could not set immutable attribute"

# Remove USB-specific markers
sudo rm -f "$TEMP_MOUNT/home/stormux/.firstboot" 2>/dev/null || true

# Install GRUB
if ! install_grub "$TEMP_MOUNT" "$TARGET_DEVICE" target_partitions; then
    error_exit "GRUB installation failed"
fi

# Validate installation
if ! validate_installation "$TEMP_MOUNT"; then
    error_exit "Installation validation failed"
fi

# Clean up all mounts
log "Unmounting filesystems..."
cleanup_mounts "$TEMP_MOUNT"

# Restore speech
restore_speech

# Success message
log "========================================="
log "Installation completed successfully!"
log "========================================="
echo
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 "Log file: $LOGFILE"
echo "========================================="
echo
echo "Press enter to continue..."
read -r