#!/bin/bash
# install-archzfs - Arch Linux ZFS Root Installation Script
# Craig Jennings <c@cjennings.net>
#
# Installs Arch Linux on ZFS root with native encryption.
# Designed to be run from the custom archzfs ISO.
#
# Features:
# - All questions asked upfront, then unattended installation
# - Optional WiFi configuration with connection test
# - ZFS native encryption (passphrase required at boot)
# - Pre-pacman ZFS snapshots for safe upgrades

set -e

#############################
# Configuration
#############################

# These will be set interactively
HOSTNAME=""
TIMEZONE=""
LOCALE="en_US.UTF-8"
KEYMAP="us"
ROOT_PASSWORD=""
ZFS_PASSPHRASE=""
WIFI_SSID=""
WIFI_PASSWORD=""

# ZFS Configuration
POOL_NAME="zroot"
COMPRESSION="zstd"
ASHIFT="12"  # 4K sectors (use 13 for 8K)

# Multi-disk RAID support
SELECTED_DISKS=()      # Array of selected disk paths (/dev/sda, /dev/sdb, ...)
ZFS_PARTS=()           # Array of ZFS partition paths
EFI_PARTS=()           # Array of EFI partition paths
RAID_LEVEL=""          # "", "mirror", "raidz1", "raidz2", "raidz3"
ENABLE_SSH="yes"       # Enable SSH with root login (default yes for headless)

# Logging
LOGFILE="/tmp/install-archzfs.log"
exec > >(tee -a "$LOGFILE") 2>&1

info()  { echo "[INFO] $1"; }
warn()  { echo "[WARN] $1"; }
error() { echo "[ERROR] $1"; exit 1; }
step()  { echo ""; echo "==> $1"; }
prompt() { echo "$1"; }

#############################
# Pre-flight Checks
#############################

preflight_checks() {
    # Check root
    [[ $EUID -ne 0 ]] && error "This script must be run as root"

    # Check ZFS module
    if ! lsmod | grep -q zfs; then
        info "Loading ZFS module..."
        modprobe zfs || error "Failed to load ZFS module. Is zfs-linux-lts installed?"
    fi

    info "ZFS module loaded successfully."
}

#############################
# Phase 1: Gather All Input
#############################

gather_input() {
    echo ""
    echo "╔═══════════════════════════════════════════════════════════════╗"
    echo "║                     Arch Linux ZFS Root                       ║"
    echo "║                Configuration and Installation                 ║"
    echo "╚═══════════════════════════════════════════════════════════════╝"
    echo ""
    info "Answer all questions now. Installation will run unattended afterward."
    echo ""

    get_hostname
    get_timezone
    get_locale
    get_keymap
    get_disks
    get_raid_level
    get_wifi
    get_zfs_passphrase
    get_root_password
    get_ssh_config
    show_summary
}

get_hostname() {
    step "Hostname"
    prompt "Enter hostname for this system:"
    read -p "> " HOSTNAME
    while [[ -z "$HOSTNAME" || ! "$HOSTNAME" =~ ^[a-zA-Z0-9]([a-zA-Z0-9-]*[a-zA-Z0-9])?$ ]]; do
        warn "Invalid hostname. Use letters, numbers, and hyphens (no spaces)."
        read -p "> " HOSTNAME
    done
}

get_timezone() {
    step "Timezone"
    prompt "Select timezone region:"
    PS3="Region: "
    select region in "America" "Europe" "Asia" "Australia" "Pacific" "Other"; do
        if [[ -n "$region" ]]; then
            break
        fi
    done

    if [[ "$region" == "Other" ]]; then
        prompt "Enter timezone (e.g., Etc/UTC):"
        read -p "> " TIMEZONE
    else
        echo ""
        prompt "Select city:"
        mapfile -t cities < <(find /usr/share/zoneinfo/"$region" -maxdepth 1 -type f -printf '%f\n' | sort)
        PS3="City: "
        select city in "${cities[@]}"; do
            if [[ -n "$city" ]]; then
                TIMEZONE="$region/$city"
                break
            fi
        done
    fi
}

get_locale() {
    step "Locale"
    prompt "Select locale:"
    PS3="Locale: "
    select loc in "en_US.UTF-8" "en_GB.UTF-8" "de_DE.UTF-8" "fr_FR.UTF-8" "es_ES.UTF-8" "Other"; do
        if [[ -n "$loc" ]]; then
            if [[ "$loc" == "Other" ]]; then
                prompt "Enter locale (e.g., ja_JP.UTF-8):"
                read -p "> " LOCALE
            else
                LOCALE="$loc"
            fi
            break
        fi
    done
}

get_keymap() {
    step "Keyboard Layout"
    prompt "Select keyboard layout:"
    PS3="Keymap: "
    select km in "us" "uk" "de" "fr" "es" "dvorak" "Other"; do
        if [[ -n "$km" ]]; then
            if [[ "$km" == "Other" ]]; then
                prompt "Enter keymap (e.g., jp106):"
                read -p "> " KEYMAP
            else
                KEYMAP="$km"
            fi
            break
        fi
    done
}

get_disks() {
    step "Disk Selection"
    echo ""
    echo "Available disks:"
    echo "----------------"
    lsblk -d -o NAME,SIZE,MODEL,TYPE | grep disk
    echo ""

    # Get list of available disks
    mapfile -t AVAILABLE_DISKS < <(lsblk -d -n -o NAME,TYPE | awk '$2=="disk"{print $1}')

    if [[ ${#AVAILABLE_DISKS[@]} -eq 0 ]]; then
        error "No disks found!"
    fi

    # Build dialog checklist items
    local dialog_items=()
    for disk in "${AVAILABLE_DISKS[@]}"; do
        local size=$(lsblk -d -n -o SIZE "/dev/$disk" | tr -d ' ')
        local model=$(lsblk -d -n -o MODEL "/dev/$disk" | tr -d ' ' | head -c 20)
        dialog_items+=("$disk" "$size $model" "off")
    done

    # Use dialog for multi-select
    local result
    result=$(dialog --stdout --checklist "Select disks for installation (SPACE to select, ENTER to confirm):" \
        20 70 10 "${dialog_items[@]}") || error "Disk selection cancelled"

    if [[ -z "$result" ]]; then
        error "No disks selected!"
    fi

    # Parse selected disks
    SELECTED_DISKS=()
    for disk in $result; do
        SELECTED_DISKS+=("/dev/$disk")
    done

    clear
    echo ""
    warn "Selected ${#SELECTED_DISKS[@]} disk(s):"
    for disk in "${SELECTED_DISKS[@]}"; do
        echo "  - $disk"
        lsblk "$disk" | sed 's/^/    /'
    done
    echo ""

    read -p "This will DESTROY all data on these disks. Type 'yes' to continue: " confirm
    if [[ "$confirm" != "yes" ]]; then
        error "Aborted by user"
    fi
}

get_raid_level() {
    local disk_count=${#SELECTED_DISKS[@]}

    if [[ $disk_count -eq 1 ]]; then
        RAID_LEVEL=""
        info "Single disk selected - no RAID"
        return
    fi

    step "RAID Configuration"
    echo ""
    echo "You have selected $disk_count disks."
    echo ""

    # Build options based on disk count
    local options=("mirror" "stripe")
    local descriptions=(
        "mirror - All disks mirror each other (max redundancy, ${disk_count}x durability)"
        "stripe - Combine disks for max capacity (NO redundancy, ${disk_count}x space)"
    )

    if [[ $disk_count -ge 3 ]]; then
        options+=("raidz1")
        descriptions+=("raidz1 - Single parity (can lose 1 disk)")
    fi
    if [[ $disk_count -ge 4 ]]; then
        options+=("raidz2")
        descriptions+=("raidz2 - Double parity (can lose 2 disks)")
    fi
    if [[ $disk_count -ge 5 ]]; then
        options+=("raidz3")
        descriptions+=("raidz3 - Triple parity (can lose 3 disks)")
    fi

    echo "Available RAID levels:"
    for i in "${!descriptions[@]}"; do
        echo "  $((i+1))) ${descriptions[$i]}"
    done
    echo ""

    PS3="Select RAID level: "
    select level in "${options[@]}"; do
        if [[ -n "$level" ]]; then
            RAID_LEVEL="$level"
            break
        fi
    done

    info "RAID level: $RAID_LEVEL"
}

get_wifi() {
    step "WiFi Configuration (Optional)"
    echo ""
    prompt "Do you want to configure WiFi? [Y/n]:"
    read -p "> " configure_wifi

    if [[ ! "$configure_wifi" =~ ^[Nn]$ ]]; then
        # Ensure NetworkManager is running
        systemctl start NetworkManager 2>/dev/null || true
        sleep 2

        echo ""
        info "Scanning for networks..."
        nmcli device wifi rescan 2>/dev/null || true
        sleep 2
        echo ""
        echo "Available networks:"
        nmcli device wifi list
        echo ""

        prompt "Enter WiFi SSID:"
        read -p "> " WIFI_SSID

        prompt "Enter WiFi password:"
        read -s -p "> " WIFI_PASSWORD
        echo ""

        # Test the connection
        info "Testing WiFi connection..."
        if nmcli device wifi connect "$WIFI_SSID" password "$WIFI_PASSWORD" 2>/dev/null; then
            info "WiFi connection successful!"
        else
            warn "WiFi connection failed. You can configure it manually after installation."
            WIFI_SSID=""
            WIFI_PASSWORD=""
        fi
    else
        info "Skipping WiFi configuration."
    fi
}

get_zfs_passphrase() {
    step "ZFS Encryption Passphrase"
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "This passphrase will be required at EVERY boot."
    echo ""
    echo "Requirements:"
    echo "  - Use a strong, memorable passphrase"
    echo "  - If forgotten, your data is UNRECOVERABLE"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    while true; do
        prompt "Enter ZFS encryption passphrase:"
        read -s -p "> " ZFS_PASSPHRASE
        echo ""

        prompt "Confirm passphrase:"
        read -s -p "> " confirm_pass
        echo ""

        if [[ "$ZFS_PASSPHRASE" == "$confirm_pass" ]]; then
            if [[ ${#ZFS_PASSPHRASE} -lt 8 ]]; then
                warn "Passphrase should be at least 8 characters."
                continue
            fi
            break
        else
            warn "Passphrases do not match. Try again."
        fi
    done
}

get_root_password() {
    step "Root Password"
    echo ""

    while true; do
        prompt "Enter root password:"
        read -s -p "> " ROOT_PASSWORD
        echo ""

        prompt "Confirm root password:"
        read -s -p "> " confirm_pass
        echo ""

        if [[ "$ROOT_PASSWORD" == "$confirm_pass" ]]; then
            break
        else
            warn "Passwords do not match. Try again."
        fi
    done
}

get_ssh_config() {
    step "SSH Configuration"
    echo ""
    info "SSH enables remote access after installation."
    info "Recommended for headless servers. Harden with archsetup later."
    echo ""
    prompt "Enable SSH with root login? [Y/n]:"
    read -p "> " ssh_choice

    if [[ "$ssh_choice" =~ ^[Nn]$ ]]; then
        ENABLE_SSH="no"
        info "SSH will not be enabled."
    else
        ENABLE_SSH="yes"
        info "SSH will be enabled with root password login."
        warn "Remember to harden SSH (key auth, fail2ban) with archsetup!"
    fi
}

show_summary() {
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Configuration Summary:"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  Hostname:     $HOSTNAME"
    echo "  Timezone:     $TIMEZONE"
    echo "  Locale:       $LOCALE"
    echo "  Keymap:       $KEYMAP"
    echo "  Disks:        ${#SELECTED_DISKS[@]} disk(s)"
    for disk in "${SELECTED_DISKS[@]}"; do
        local size=$(lsblk -d -n -o SIZE "$disk" | tr -d ' ')
        echo "                - $disk ($size)"
    done
    echo "  RAID Level:   ${RAID_LEVEL:-single (no RAID)}"
    echo "  WiFi:         ${WIFI_SSID:-Not configured}"
    echo "  SSH:          ${ENABLE_SSH:-yes} (root login)"
    echo "  ZFS Pool:     $POOL_NAME (encrypted)"
    echo "  Boot:         EFI on all disks (redundant)"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    read -p "Press Enter to begin installation, or Ctrl+C to abort..."
}

#############################
# Phase 2: Installation
#############################

partition_disks() {
    step "Partitioning ${#SELECTED_DISKS[@]} disk(s)"

    EFI_PARTS=()
    ZFS_PARTS=()

    for disk in "${SELECTED_DISKS[@]}"; do
        info "Partitioning $disk..."

        # Wipe existing signatures
        wipefs -af "$disk"
        sgdisk --zap-all "$disk"

        # Create partitions: 1G EFI + rest for ZFS
        sgdisk -n 1:0:+1G -t 1:ef00 -c 1:"EFI" "$disk"
        sgdisk -n 2:0:0 -t 2:bf00 -c 2:"ZFS" "$disk"

        # Determine partition names (handle nvme/mmcblk naming)
        local efi_part zfs_part
        if [[ "$disk" == *"nvme"* ]] || [[ "$disk" == *"mmcblk"* ]]; then
            efi_part="${disk}p1"
            zfs_part="${disk}p2"
        else
            efi_part="${disk}1"
            zfs_part="${disk}2"
        fi

        EFI_PARTS+=("$efi_part")
        ZFS_PARTS+=("$zfs_part")

        sleep 1
        partprobe "$disk"
    done

    sleep 2

    # Format all EFI partitions
    for i in "${!EFI_PARTS[@]}"; do
        info "Formatting EFI partition ${EFI_PARTS[$i]}..."
        mkfs.fat -F32 -n "EFI$i" "${EFI_PARTS[$i]}"
    done

    info "Partitioning complete. Created ${#EFI_PARTS[@]} EFI and ${#ZFS_PARTS[@]} ZFS partitions."
}

create_zfs_pool() {
    step "Creating ZFS Pool with Native Encryption"

    if zpool list "$POOL_NAME" &>/dev/null; then
        warn "Pool $POOL_NAME already exists. Destroying..."
        zpool destroy -f "$POOL_NAME"
    fi

    # Build pool configuration based on RAID level
    local pool_config
    if [[ "$RAID_LEVEL" == "stripe" ]]; then
        # Stripe: just list devices without a vdev type (RAID0 equivalent)
        pool_config="${ZFS_PARTS[*]}"
        info "Creating striped pool with ${#ZFS_PARTS[@]} disks (NO redundancy)..."
        warn "Data loss will occur if ANY disk fails!"
    elif [[ -n "$RAID_LEVEL" ]]; then
        pool_config="$RAID_LEVEL ${ZFS_PARTS[*]}"
        info "Creating $RAID_LEVEL pool with ${#ZFS_PARTS[@]} disks..."
    else
        pool_config="${ZFS_PARTS[0]}"
        info "Creating single-disk pool..."
    fi

    # Create encrypted pool using passphrase from variable
    echo "$ZFS_PASSPHRASE" | zpool create -f \
        -o ashift="$ASHIFT" \
        -o autotrim=on \
        -O acltype=posixacl \
        -O atime=off \
        -O canmount=off \
        -O compression="$COMPRESSION" \
        -O dnodesize=auto \
        -O normalization=formD \
        -O relatime=on \
        -O xattr=sa \
        -O encryption=aes-256-gcm \
        -O keyformat=passphrase \
        -O keylocation=prompt \
        -O mountpoint=none \
        -R /mnt \
        "$POOL_NAME" $pool_config

    info "ZFS pool created successfully."
    zpool status "$POOL_NAME"
}

create_datasets() {
    step "Creating ZFS Datasets"

    # Root dataset container
    zfs create -o mountpoint=none -o canmount=off "$POOL_NAME/ROOT"

    # Main root filesystem
    # Reserve 20% of pool or 20G max to prevent pool from filling completely
    local pool_size_bytes=$(zpool get -Hp size "$POOL_NAME" | awk '{print $3}')
    local pool_size_gb=$((pool_size_bytes / 1024 / 1024 / 1024))
    local reserve_gb=$((pool_size_gb / 5))  # 20%
    [[ $reserve_gb -gt 20 ]] && reserve_gb=20
    [[ $reserve_gb -lt 5 ]] && reserve_gb=5

    zfs create -o mountpoint=/ -o canmount=noauto -o reservation=${reserve_gb}G "$POOL_NAME/ROOT/default"
    zfs mount "$POOL_NAME/ROOT/default"

    # Home (archsetup will create user subdataset)
    zfs create -o mountpoint=/home "$POOL_NAME/home"
    zfs create -o mountpoint=/root "$POOL_NAME/home/root"

    # Media - compression off for already-compressed files
    zfs create -o mountpoint=/media -o compression=off "$POOL_NAME/media"

    # VMs - 64K recordsize for VM disk images
    zfs create -o mountpoint=/vms -o recordsize=64K "$POOL_NAME/vms"

    # Var datasets
    zfs create -o mountpoint=/var -o canmount=off "$POOL_NAME/var"
    zfs create -o mountpoint=/var/log "$POOL_NAME/var/log"
    zfs create -o mountpoint=/var/cache "$POOL_NAME/var/cache"
    zfs create -o mountpoint=/var/lib -o canmount=off "$POOL_NAME/var/lib"
    zfs create -o mountpoint=/var/lib/pacman "$POOL_NAME/var/lib/pacman"
    zfs create -o mountpoint=/var/lib/docker "$POOL_NAME/var/lib/docker"

    # Temp directories - excluded from snapshots
    zfs create -o mountpoint=/var/tmp -o com.sun:auto-snapshot=false "$POOL_NAME/var/tmp"
    zfs create -o mountpoint=/tmp -o com.sun:auto-snapshot=false "$POOL_NAME/tmp"
    chmod 1777 /mnt/tmp /mnt/var/tmp

    info "Datasets created:"
    zfs list -r "$POOL_NAME" -o name,mountpoint,compression
}

mount_efi() {
    step "Mounting EFI Partition"
    mkdir -p /mnt/boot
    # Mount primary (first) EFI partition
    mount "${EFI_PARTS[0]}" /mnt/boot
    info "Primary EFI partition ${EFI_PARTS[0]} mounted at /mnt/boot"
}

install_base() {
    step "Installing Base System"

    info "Updating pacman keys..."
    pacman-key --init
    pacman-key --populate archlinux

    # Add archzfs key
    pacman-key -r DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true
    pacman-key --lsign-key DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true

    # Add archzfs repo to pacman.conf for pacstrap
    if ! grep -q "\[archzfs\]" /etc/pacman.conf; then
        cat >> /etc/pacman.conf << 'EOF'

[archzfs]
Server = https://archzfs.com/$repo/$arch
SigLevel = Optional TrustAll
EOF
    fi

    info "Installing base packages (this takes a while)..."
    info "ZFS will be built from source via DKMS - this ensures kernel compatibility."
    pacstrap -K /mnt \
        base \
        base-devel \
        linux-lts \
        linux-lts-headers \
        linux-firmware \
        zfs-dkms \
        zfs-utils \
        grub \
        freetype2 \
        efibootmgr \
        networkmanager \
        openssh \
        git \
        vim \
        sudo \
        zsh \
        nodejs \
        npm \
        ttf-dejavu

    info "Base system installed."
}

configure_system() {
    step "Configuring System"

    # fstab (only for EFI)
    info "Generating fstab..."
    echo "# /boot - EFI System Partition" > /mnt/etc/fstab
    echo "UUID=$(blkid -s UUID -o value "$EFI_PART")  /boot  vfat  defaults,noatime  0 2" >> /mnt/etc/fstab

    # Timezone
    info "Setting timezone to $TIMEZONE..."
    arch-chroot /mnt ln -sf "/usr/share/zoneinfo/$TIMEZONE" /etc/localtime
    arch-chroot /mnt hwclock --systohc

    # Locale
    info "Configuring locale..."
    echo "$LOCALE UTF-8" >> /mnt/etc/locale.gen
    arch-chroot /mnt locale-gen
    echo "LANG=$LOCALE" > /mnt/etc/locale.conf

    # Keymap
    echo "KEYMAP=$KEYMAP" > /mnt/etc/vconsole.conf

    # Hostname
    info "Setting hostname to $HOSTNAME..."
    echo "$HOSTNAME" > /mnt/etc/hostname
    cat > /mnt/etc/hosts << EOF
127.0.0.1   localhost
::1         localhost
127.0.1.1   $HOSTNAME.localdomain $HOSTNAME
EOF

    # Add archzfs repo
    info "Adding archzfs repository..."
    cat >> /mnt/etc/pacman.conf << 'EOF'

[archzfs]
Server = https://archzfs.com/$repo/$arch
SigLevel = Optional TrustAll
EOF

    # Import archzfs key
    arch-chroot /mnt pacman-key -r DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true
    arch-chroot /mnt pacman-key --lsign-key DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true

    # Set root password
    info "Setting root password..."
    echo "root:$ROOT_PASSWORD" | arch-chroot /mnt chpasswd
}

configure_wifi() {
    if [[ -n "$WIFI_SSID" ]]; then
        step "Configuring WiFi"

        # Copy NetworkManager connection from live environment
        if [[ -d /etc/NetworkManager/system-connections ]]; then
            mkdir -p /mnt/etc/NetworkManager/system-connections
            cp /etc/NetworkManager/system-connections/* /mnt/etc/NetworkManager/system-connections/ 2>/dev/null || true
            chmod 600 /mnt/etc/NetworkManager/system-connections/* 2>/dev/null || true
        fi

        info "WiFi configuration copied to installed system."
    fi
}

configure_ssh() {
    if [[ "$ENABLE_SSH" == "yes" ]]; then
        step "Configuring SSH"

        # Ensure sshd config allows root login with password
        sed -i 's/^#PermitRootLogin.*/PermitRootLogin yes/' /mnt/etc/ssh/sshd_config
        sed -i 's/^PermitRootLogin.*/PermitRootLogin yes/' /mnt/etc/ssh/sshd_config

        # Enable sshd service
        arch-chroot /mnt systemctl enable sshd

        info "SSH enabled with root password login."
        warn "Run archsetup to harden SSH (key auth, fail2ban)."
    else
        info "SSH not enabled. Enable manually if needed."
    fi
}

configure_initramfs() {
    step "Configuring Initramfs for ZFS"

    cp /mnt/etc/mkinitcpio.conf /mnt/etc/mkinitcpio.conf.bak

    # Configure hooks for ZFS
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect microcode modconf kms keyboard keymap consolefont block zfs filesystems fsck)/' /mnt/etc/mkinitcpio.conf

    info "Regenerating initramfs..."
    arch-chroot /mnt mkinitcpio -P
}

configure_bootloader() {
    step "Configuring GRUB Bootloader"

    # Configure GRUB defaults
    cat > /mnt/etc/default/grub << EOF
GRUB_DEFAULT=0
GRUB_TIMEOUT=5
GRUB_DISTRIBUTOR="Arch Linux (ZFS)"
GRUB_CMDLINE_LINUX_DEFAULT="loglevel=3 quiet"
GRUB_CMDLINE_LINUX="root=ZFS=$POOL_NAME/ROOT/default"
GRUB_PRELOAD_MODULES="part_gpt part_msdos zfs"
GRUB_TERMINAL_OUTPUT="console"
GRUB_DISABLE_OS_PROBER=true
GRUB_GFXMODE=auto
GRUB_GFXPAYLOAD_LINUX=keep
GRUB_FONT=/boot/grub/fonts/DejaVuSansMono32.pf2
EOF

    # Install GRUB to each EFI partition for boot redundancy
    info "Installing GRUB to ${#EFI_PARTS[@]} EFI partition(s)..."

    for i in "${!EFI_PARTS[@]}"; do
        local efi_part="${EFI_PARTS[$i]}"
        local bootloader_id="GRUB"
        if [[ ${#EFI_PARTS[@]} -gt 1 ]]; then
            bootloader_id="GRUB-disk$((i+1))"
        fi

        # Unmount current boot if mounted, mount this EFI partition
        umount /mnt/boot 2>/dev/null || true
        mount "$efi_part" /mnt/boot

        # Create directories and font
        mkdir -p /mnt/boot/grub/fonts
        arch-chroot /mnt grub-mkfont -s 32 -o /boot/grub/fonts/DejaVuSansMono32.pf2 \
            /usr/share/fonts/TTF/DejaVuSansMono.ttf 2>/dev/null || true

        # Install GRUB
        info "Installing GRUB to $efi_part (bootloader-id: $bootloader_id)..."
        arch-chroot /mnt grub-install --target=x86_64-efi --efi-directory=/boot \
            --bootloader-id="$bootloader_id" --recheck

        # Generate configuration
        arch-chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
    done

    # Remount primary EFI for rest of installation
    umount /mnt/boot 2>/dev/null || true
    mount "${EFI_PARTS[0]}" /mnt/boot

    if [[ ${#EFI_PARTS[@]} -gt 1 ]]; then
        info "GRUB installed to all ${#EFI_PARTS[@]} disks for boot redundancy."
    fi
}

configure_zfs_services() {
    step "Configuring ZFS Services"

    arch-chroot /mnt systemctl enable zfs.target
    arch-chroot /mnt systemctl enable zfs-import-cache
    arch-chroot /mnt systemctl enable zfs-mount
    arch-chroot /mnt systemctl enable zfs-import.target

    # Generate zpool cache
    mkdir -p /mnt/etc/zfs
    zpool set cachefile=/etc/zfs/zpool.cache "$POOL_NAME"
    cp /etc/zfs/zpool.cache /mnt/etc/zfs/

    # Set bootfs
    zpool set bootfs="$POOL_NAME/ROOT/default" "$POOL_NAME"

    # Enable other services
    arch-chroot /mnt systemctl enable NetworkManager
    arch-chroot /mnt systemctl enable sshd

    info "ZFS services configured."
}

configure_pacman_hook() {
    step "Configuring Pacman Snapshot Hook"

    mkdir -p /mnt/etc/pacman.d/hooks

    cat > /mnt/etc/pacman.d/hooks/zfs-snapshot.hook << EOF
[Trigger]
Operation = Upgrade
Operation = Install
Operation = Remove
Type = Package
Target = *

[Action]
Description = Creating ZFS snapshot before pacman transaction...
When = PreTransaction
Exec = /usr/local/bin/zfs-pre-snapshot
EOF

    cat > /mnt/usr/local/bin/zfs-pre-snapshot << 'EOF'
#!/bin/bash
POOL="zroot"
DATASET="$POOL/ROOT/default"
TIMESTAMP=$(date +%Y-%m-%d_%H-%M-%S)
SNAPSHOT_NAME="pre-pacman_$TIMESTAMP"

if zfs snapshot "$DATASET@$SNAPSHOT_NAME"; then
    echo "Created snapshot: $DATASET@$SNAPSHOT_NAME"
else
    echo "Warning: Failed to create snapshot" >&2
fi
EOF

    chmod +x /mnt/usr/local/bin/zfs-pre-snapshot

    info "Pacman hook configured."
}

copy_archsetup() {
    step "Installing archsetup Launcher"

    cat > /mnt/usr/local/bin/archsetup << 'EOF'
#!/bin/bash
curl -fsSL https://cjennings.net/archsetup | bash
EOF
    chmod +x /mnt/usr/local/bin/archsetup
    info "archsetup launcher installed to /usr/local/bin/archsetup"
}

sync_efi_partitions() {
    # Skip if only one disk
    if [[ ${#EFI_PARTS[@]} -le 1 ]]; then
        return
    fi

    step "Syncing EFI Partitions for Redundancy"

    local primary_efi="${EFI_PARTS[0]}"
    local temp_mount="/mnt/efi_sync"

    for i in "${!EFI_PARTS[@]}"; do
        if [[ $i -eq 0 ]]; then
            continue  # Skip primary
        fi

        local efi_part="${EFI_PARTS[$i]}"
        info "Syncing to EFI partition $((i+1)): $efi_part"

        mkdir -p "$temp_mount"
        mount "$efi_part" "$temp_mount"

        # Sync all content from primary EFI (mounted at /mnt/boot) to secondary
        rsync -a --delete /mnt/boot/ "$temp_mount/"

        umount "$temp_mount"
    done

    rmdir "$temp_mount" 2>/dev/null || true
    info "All EFI partitions synchronized."
}

create_genesis_snapshot() {
    step "Creating Genesis Snapshot"

    # Create recursive snapshot of entire pool
    info "Creating snapshot ${POOL_NAME}@genesis..."
    zfs snapshot -r "${POOL_NAME}@genesis"

    # Create rollback script in /root
    info "Installing rollback-to-genesis script..."
    cat > /mnt/root/rollback-to-genesis << 'ROLLBACK_EOF'
#!/bin/bash
# rollback-to-genesis - Roll back all datasets to the genesis snapshot
#
# This script rolls back the entire ZFS pool to its pristine post-install state.
# WARNING: This will destroy all changes made since installation!

set -e

POOL_NAME="zroot"

echo "╔═══════════════════════════════════════════════════════════════╗"
echo "║              WARNING: Full System Rollback                    ║"
echo "╚═══════════════════════════════════════════════════════════════╝"
echo ""
echo "This will roll back ALL datasets to the genesis snapshot!"
echo "All changes since installation will be permanently lost."
echo ""

# Show what will be rolled back
echo "Datasets to roll back:"
zfs list -r -t snapshot -o name "${POOL_NAME}" 2>/dev/null | grep "@genesis" | while read snap; do
    dataset="${snap%@genesis}"
    echo "  - $dataset"
done
echo ""

read -p "Type 'ROLLBACK' to confirm: " confirm
if [[ "$confirm" != "ROLLBACK" ]]; then
    echo "Aborted."
    exit 1
fi

echo ""
echo "Rolling back to genesis..."

# Roll back each dataset (must do in reverse order for dependencies)
zfs list -r -H -o name "${POOL_NAME}" | tac | while read dataset; do
    if zfs list -t snapshot "${dataset}@genesis" &>/dev/null; then
        echo "  Rolling back: $dataset"
        zfs rollback -r "${dataset}@genesis"
    fi
done

echo ""
echo "Rollback complete!"
echo "Reboot to complete the process: reboot"
ROLLBACK_EOF

    chmod +x /mnt/root/rollback-to-genesis
    info "Genesis snapshot created. Rollback script: /root/rollback-to-genesis"
}

cleanup() {
    step "Cleaning Up"

    # Clear sensitive variables
    ROOT_PASSWORD=""
    ZFS_PASSPHRASE=""

    info "Unmounting filesystems..."
    umount /mnt/boot 2>/dev/null || true

    info "Exporting ZFS pool..."
    zpool export "$POOL_NAME"

    info "Cleanup complete."
}

print_summary() {
    echo ""
    echo "╔═══════════════════════════════════════════════════════════════╗"
    echo "║                    Installation Complete!                     ║"
    echo "╚═══════════════════════════════════════════════════════════════╝"
    echo ""
    echo "System Configuration:"
    echo "  Hostname:     $HOSTNAME"
    echo "  Timezone:     $TIMEZONE"
    echo "  ZFS Pool:     $POOL_NAME (encrypted)"
    echo ""
    echo "ZFS Features:"
    echo "  - Genesis snapshot: pristine post-install state"
    echo "  - Pre-pacman snapshots for safe upgrades"
    echo "  - Sanoid/syncoid configured by archsetup"
    echo ""
    echo "Next Steps:"
    echo "  1. Reboot: reboot"
    echo "  2. Enter ZFS encryption passphrase at boot"
    echo "  3. Log in as root"
    echo "  4. Run archsetup: archsetup"
    echo ""
    echo "Useful Commands:"
    echo "  List snapshots:    zfs list -t snapshot"
    echo "  Manual snapshot:   zfs snapshot zroot/home@my-backup"
    echo "  Rollback:          zfs rollback zroot/home@my-backup"
    echo "  Factory reset:     /root/rollback-to-genesis"
    echo "  Pool status:       zpool status"
    echo ""
    info "Installation log: $LOGFILE"
    echo ""
}

#############################
# Main
#############################

main() {
    preflight_checks
    gather_input

    # Unattended installation begins
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Beginning unattended installation..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    partition_disks
    create_zfs_pool
    create_datasets
    mount_efi
    install_base
    configure_system
    configure_wifi
    configure_ssh
    configure_initramfs
    configure_bootloader
    configure_zfs_services
    configure_pacman_hook
    copy_archsetup
    sync_efi_partitions
    create_genesis_snapshot
    cleanup
    print_summary
}

trap 'error "Installation interrupted!"' INT TERM

main "$@"