#!/bin/bash
# archangel - Arch Linux Installer with Snapshot-Based Recovery
# Craig Jennings <c@cjennings.net>
#
# Installs Arch Linux on ZFS or Btrfs root with snapshot support.
# Choose your filesystem: ZFS (native encryption) or Btrfs (GRUB snapshots).
#
# Features:
# - Filesystem choice: ZFS or Btrfs
# - All questions asked upfront, then unattended installation
# - Optional WiFi configuration with connection test
# - Optional ZFS native encryption (passphrase required at boot)
# - Pre-pacman snapshots for safe upgrades
# - Genesis snapshot for factory reset
#
# UNATTENDED MODE:
#   Use --config-file /path/to/archangel.conf for automated installs.
#   Config file must be explicitly specified to prevent accidental disk wipes.
#   See /root/archangel.conf.example for a template with all options.

set -e

#############################
# Source Library Functions
#############################

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/lib/common.sh"
source "$SCRIPT_DIR/lib/config.sh"
source "$SCRIPT_DIR/lib/disk.sh"
source "$SCRIPT_DIR/lib/zfs.sh"
source "$SCRIPT_DIR/lib/btrfs.sh"

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

# Filesystem selection (zfs or btrfs)
FILESYSTEM="zfs"       # Default to ZFS, can be changed interactively or via config

# 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)
NO_ENCRYPT="no"        # Skip ZFS encryption (for testing only)

# Logging
LOGFILE="/tmp/archangel-$(date +'%Y-%m-%d-%H-%M-%S').log"
exec > >(tee -a "$LOGFILE") 2>&1

# Log header with timestamp
echo ""
echo "================================================================================"
echo "archangel started @ $(date +'%Y-%m-%d %H:%M:%S')"
echo "================================================================================"
echo ""

# Output functions now in lib/common.sh
# Config functions now in lib/config.sh

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

preflight_checks() {
    require_root
}

# Filesystem-specific preflight (called after filesystem is selected)
filesystem_preflight() {
    if [[ "$FILESYSTEM" == "zfs" ]]; then
        zfs_preflight
    elif [[ "$FILESYSTEM" == "btrfs" ]]; then
        btrfs_preflight
    fi
}

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

gather_input() {
    if [[ "$UNATTENDED" == true ]]; then
        # Validate required config values
        if [[ -z "$HOSTNAME" ]]; then error "Config missing required: HOSTNAME"; fi
        if [[ -z "$TIMEZONE" ]]; then error "Config missing required: TIMEZONE"; fi
        if [[ -z "$ROOT_PASSWORD" ]]; then error "Config missing required: ROOT_PASSWORD"; fi
        if [[ ${#SELECTED_DISKS[@]} -eq 0 ]]; then error "Config missing required: DISKS"; fi

        # Set defaults for optional values
        [[ -z "$FILESYSTEM" ]] && FILESYSTEM="zfs" || true
        [[ -z "$LOCALE" ]] && LOCALE="en_US.UTF-8" || true
        [[ -z "$KEYMAP" ]] && KEYMAP="us" || true
        [[ -z "$ENABLE_SSH" ]] && ENABLE_SSH="yes" || true

        # ZFS-specific validation
        if [[ "$FILESYSTEM" == "zfs" ]]; then
            if [[ "$NO_ENCRYPT" != "yes" && -z "$ZFS_PASSPHRASE" ]]; then
                error "Config missing required: ZFS_PASSPHRASE (or set NO_ENCRYPT=yes)"
            fi
        fi

        # Btrfs-specific validation
        if [[ "$FILESYSTEM" == "btrfs" ]]; then
            if [[ "$NO_ENCRYPT" != "yes" && -z "$LUKS_PASSPHRASE" ]]; then
                error "Config missing required: LUKS_PASSPHRASE (or set NO_ENCRYPT=yes)"
            fi
        fi

        # Validate filesystem choice
        if [[ "$FILESYSTEM" != "zfs" && "$FILESYSTEM" != "btrfs" ]]; then
            error "Invalid FILESYSTEM: $FILESYSTEM (must be 'zfs' or 'btrfs')"
        fi

        # Determine RAID level if not specified
        if [[ -z "$RAID_LEVEL" && ${#SELECTED_DISKS[@]} -gt 1 ]]; then
            RAID_LEVEL="mirror"
            info "Defaulting to mirror for ${#SELECTED_DISKS[@]} disks"
        fi

        info "Configuration loaded:"
        info "  Filesystem: $FILESYSTEM"
        info "  Hostname: $HOSTNAME"
        info "  Timezone: $TIMEZONE"
        info "  Locale: $LOCALE"
        info "  Keymap: $KEYMAP"
        info "  Disks: ${SELECTED_DISKS[*]}"
        [[ -n "$RAID_LEVEL" ]] && info "  RAID: $RAID_LEVEL"
        info "  SSH: $ENABLE_SSH"
        [[ "$NO_ENCRYPT" == "yes" ]] && warn "  Encryption: DISABLED (testing mode)"
        [[ -n "$WIFI_SSID" ]] && info "  WiFi: $WIFI_SSID"
        return 0
    fi

    echo ""
    echo "╔═══════════════════════════════════════════════════════════════╗"
    echo "║                        Archangel                              ║"
    echo "║            Arch Linux with Snapshot-Based Recovery            ║"
    echo "╚═══════════════════════════════════════════════════════════════╝"
    echo ""
    info "Answer all questions now. Installation will run unattended afterward."
    echo ""

    select_filesystem
    get_hostname
    get_timezone
    get_locale
    get_keymap
    get_disks
    get_raid_level
    get_wifi

    # Encryption handling (filesystem-specific)
    if [[ "$FILESYSTEM" == "zfs" ]]; then
        get_encryption_choice
        [[ "$NO_ENCRYPT" != "yes" ]] && get_zfs_passphrase
    elif [[ "$FILESYSTEM" == "btrfs" ]]; then
        get_btrfs_encryption_choice
        [[ "$NO_ENCRYPT" != "yes" ]] && get_luks_passphrase
    fi

    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"
    echo ""
    info "Type to search, ENTER to select"
    echo ""

    TIMEZONE=$(find /usr/share/zoneinfo -type f ! -path '*/posix/*' ! -path '*/right/*' \
        | sed 's|/usr/share/zoneinfo/||' \
        | sort \
        | fzf --height=20 --layout=reverse --border \
              --header="Select Timezone" \
              --preview='echo "Timezone: {}"; echo ""; TZ={} date "+Current time: %Y-%m-%d %H:%M:%S %Z"' \
              --preview-window=right:40%)

    if [[ -z "$TIMEZONE" ]]; then
        error "No timezone selected!"
    fi
    info "Selected: $TIMEZONE"
}

get_locale() {
    step "Locale"
    echo ""
    info "Type to search, ENTER to select"
    echo ""

    # Get available locales from locale.gen
    LOCALE=$(grep -E "^#?[a-z]" /etc/locale.gen \
        | sed 's/^#//' \
        | awk '{print $1}' \
        | sort -u \
        | fzf --height=20 --layout=reverse --border \
              --header="Select Locale (type to search, e.g. 'de_DE', 'fr_FR')" \
              --preview='
                  loc={}
                  echo "Locale: $loc"
                  echo ""
                  lang=${loc%%_*}
                  country=${loc#*_}
                  country=${country%%.*}
                  echo "Language: $lang"
                  echo "Country:  $country"
                  echo ""
                  echo "Example formats:"
                  echo "  Date:     $(LC_ALL={} date "+%x" 2>/dev/null || echo "N/A")"
                  echo "  Currency: $(LC_ALL={} locale currency_symbol 2>/dev/null || echo "N/A")"
              ' \
              --preview-window=right:45%)

    if [[ -z "$LOCALE" ]]; then
        error "No locale selected!"
    fi
    info "Selected: $LOCALE"
}

get_keymap() {
    step "Keyboard Layout"
    echo ""
    info "Type to search, ENTER to select"
    echo ""

    KEYMAP=$(localectl list-keymaps \
        | fzf --height=20 --layout=reverse --border \
              --header="Select Keyboard Layout (type to search)" \
              --preview='
                  echo "Keymap: {}"
                  echo ""
                  echo "This will set your console keyboard layout."
                  echo ""
                  echo "Common layouts:"
                  echo "  us      - US English (QWERTY)"
                  echo "  uk      - UK English"
                  echo "  de      - German (QWERTZ)"
                  echo "  fr      - French (AZERTY)"
                  echo "  dvorak  - Dvorak"
              ' \
              --preview-window=right:45%)

    if [[ -z "$KEYMAP" ]]; then
        error "No keymap selected!"
    fi
    info "Selected: $KEYMAP"
}

get_disks() {
    step "Disk Selection"
    echo ""
    info "TAB to select multiple disks, ENTER to confirm"
    echo ""

    # Get list of available disks with info
    local disk_list
    disk_list=$(lsblk -d -n -o NAME,SIZE,TYPE | awk '$3=="disk"{printf "/dev/%-8s %8s\n", $1, $2}')

    if [[ -z "$disk_list" ]]; then
        error "No disks found!"
    fi

    # Use fzf for multi-select with disk details preview
    local selected
    selected=$(echo "$disk_list" \
        | fzf --multi --height=20 --layout=reverse --border \
              --header="Select Disks (TAB to toggle, ENTER to confirm)" \
              --preview='
                  disk=$(echo {} | awk "{print \$1}")
                  echo "Disk: $disk"
                  echo ""
                  echo "Details:"
                  lsblk -o NAME,SIZE,TYPE,FSTYPE,MOUNTPOINT "$disk" 2>/dev/null
                  echo ""
                  echo "Disk info:"
                  udevadm info --query=property "$disk" 2>/dev/null | grep -E "ID_MODEL=|ID_SERIAL=" | sed "s/^/  /"
              ' \
              --preview-window=right:50%)

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

    # Parse selected disks
    SELECTED_DISKS=()
    while IFS= read -r line; do
        local disk=$(echo "$line" | awk '{print $1}')
        SELECTED_DISKS+=("$disk")
    done <<< "$selected"

    echo ""
    warn "Selected ${#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 ""

    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 ""
    info "Select RAID level (ENTER to confirm)"
    echo ""

    # Calculate total raw size for preview
    local total_bytes=0
    local smallest_bytes=0
    for disk in "${SELECTED_DISKS[@]}"; do
        local bytes=$(lsblk -b -d -n -o SIZE "$disk")
        total_bytes=$((total_bytes + bytes))
        if [[ $smallest_bytes -eq 0 ]] || [[ $bytes -lt $smallest_bytes ]]; then
            smallest_bytes=$bytes
        fi
    done
    local total_gb=$((total_bytes / 1073741824))
    local smallest_gb=$((smallest_bytes / 1073741824))

    # Build options based on disk count
    local options="mirror\nstripe"
    [[ $disk_count -ge 3 ]] && options+="\nraidz1"
    [[ $disk_count -ge 4 ]] && options+="\nraidz2"
    [[ $disk_count -ge 5 ]] && options+="\nraidz3"

    # Export variables for preview subshell
    export RAID_DISK_COUNT=$disk_count
    export RAID_TOTAL_GB=$total_gb
    export RAID_SMALLEST_GB=$smallest_gb

    RAID_LEVEL=$(echo -e "$options" \
        | fzf --height=20 --layout=reverse --border \
              --header="Select RAID Level ($disk_count disks, ${total_gb}GB total)" \
              --preview='
                  n=$RAID_DISK_COUNT
                  total=$RAID_TOTAL_GB
                  small=$RAID_SMALLEST_GB

                  case {} in
                    mirror)
                      echo "MIRROR"
                      echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                      echo ""
                      echo "All disks contain identical copies of data."
                      echo "Maximum redundancy - can survive loss of"
                      echo "all disks except one."
                      echo ""
                      echo "Redundancy:    Can lose $((n-1)) of $n disks"
                      echo "Usable space:  ~${small}GB (smallest disk)"
                      echo "Read speed:    Fast (parallel reads)"
                      echo "Write speed:   Normal"
                      echo ""
                      echo "Best for:"
                      echo "  - Boot drives"
                      echo "  - Critical data"
                      echo "  - Maximum safety"
                      ;;
                    stripe)
                      echo "STRIPE (RAID0)"
                      echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                      echo ""
                      echo "WARNING: NO REDUNDANCY!"
                      echo "Data is striped across all disks."
                      echo "ANY disk failure = ALL data lost!"
                      echo ""
                      echo "Redundancy:    NONE"
                      echo "Usable space:  ~${total}GB (all disks)"
                      echo "Read speed:    Very fast"
                      echo "Write speed:   Very fast"
                      echo ""
                      echo "Best for:"
                      echo "  - Scratch/temp space"
                      echo "  - Replaceable data"
                      echo "  - Maximum performance"
                      ;;
                    raidz1)
                      usable=$(( (n-1) * small ))
                      echo "RAIDZ1 (Single Parity)"
                      echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                      echo ""
                      echo "One disk worth of parity distributed"
                      echo "across all disks."
                      echo ""
                      echo "Redundancy:    Can lose 1 disk"
                      echo "Usable space:  ~${usable}GB ($((n-1)) of $n disks)"
                      echo "Read speed:    Fast"
                      echo "Write speed:   Good"
                      echo ""
                      echo "Best for:"
                      echo "  - General storage"
                      echo "  - Good balance of space/safety"
                      ;;
                    raidz2)
                      usable=$(( (n-2) * small ))
                      echo "RAIDZ2 (Double Parity)"
                      echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                      echo ""
                      echo "Two disks worth of parity distributed"
                      echo "across all disks."
                      echo ""
                      echo "Redundancy:    Can lose 2 disks"
                      echo "Usable space:  ~${usable}GB ($((n-2)) of $n disks)"
                      echo "Read speed:    Fast"
                      echo "Write speed:   Good"
                      echo ""
                      echo "Best for:"
                      echo "  - Large arrays (5+ disks)"
                      echo "  - Important data"
                      ;;
                    raidz3)
                      usable=$(( (n-3) * small ))
                      echo "RAIDZ3 (Triple Parity)"
                      echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                      echo ""
                      echo "Three disks worth of parity distributed"
                      echo "across all disks."
                      echo ""
                      echo "Redundancy:    Can lose 3 disks"
                      echo "Usable space:  ~${usable}GB ($((n-3)) of $n disks)"
                      echo "Read speed:    Fast"
                      echo "Write speed:   Moderate"
                      echo ""
                      echo "Best for:"
                      echo "  - Very large arrays (8+ disks)"
                      echo "  - Archival storage"
                      ;;
                  esac
              ' \
              --preview-window=right:50%)

    # Clean up exported variables
    unset RAID_DISK_COUNT RAID_TOTAL_GB RAID_SMALLEST_GB

    if [[ -z "$RAID_LEVEL" ]]; then
        error "No RAID level selected!"
    fi
    info "Selected: $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 3

        # Get list of networks for fzf
        local networks
        networks=$(nmcli -t -f SSID,SIGNAL,SECURITY device wifi list | grep -v '^$' | sort -t: -k2 -rn | uniq)

        if [[ -z "$networks" ]]; then
            warn "No WiFi networks found."
            info "Skipping WiFi configuration."
            return
        fi

        echo ""
        info "Select network (ENTER to confirm, ESC to skip)"
        echo ""

        # Use fzf to select network
        WIFI_SSID=$(echo "$networks" \
            | fzf --height=15 --layout=reverse --border \
                  --header="Select WiFi Network" \
                  --delimiter=':' \
                  --with-nth=1 \
                  --preview='
                      IFS=":" read -r ssid signal security <<< {}
                      echo "Network: $ssid"
                      echo ""
                      echo "Signal:   ${signal}%"
                      echo "Security: ${security:-Open}"
                      echo ""
                      if [[ -z "$security" ]]; then
                          echo "WARNING: Open network (no encryption)"
                      fi
                  ' \
                  --preview-window=right:40% \
            | cut -d: -f1)

        if [[ -z "$WIFI_SSID" ]]; then
            info "Skipping WiFi configuration."
            return
        fi

        prompt "Enter WiFi password for '$WIFI_SSID':"
        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_btrfs_encryption_choice() {
    step "Btrfs Encryption (LUKS)"

    echo ""
    echo "LUKS encryption protects your data at rest."
    echo "You'll need to enter a passphrase at each boot."
    echo ""

    prompt "Enable LUKS encryption? [Y/n]:"
    read -p "> " encrypt_choice

    if [[ "$encrypt_choice" =~ ^[Nn]$ ]]; then
        NO_ENCRYPT="yes"
        warn "Encryption DISABLED - data will not be encrypted at rest"
    else
        NO_ENCRYPT="no"
        info "LUKS encryption enabled - you'll set a passphrase next"
    fi
}

get_luks_passphrase() {
    step "LUKS Encryption Passphrase"

    echo ""
    echo "Choose a strong passphrase for disk encryption."
    echo "You'll need this passphrase every time you boot."
    echo ""
    echo "IMPORTANT: If you forget this passphrase, your data is UNRECOVERABLE!"
    echo ""

    while true; do
        prompt "Enter LUKS encryption passphrase:"
        read -rs LUKS_PASSPHRASE
        echo ""

        prompt "Confirm passphrase:"
        read -rs confirm
        echo ""

        if [[ "$LUKS_PASSPHRASE" == "$confirm" ]]; then
            if [[ ${#LUKS_PASSPHRASE} -lt 8 ]]; then
                warn "Passphrase should be at least 8 characters. Try again."
            else
                info "Passphrase confirmed."
                break
            fi
        else
            warn "Passphrases don't match. Try again."
        fi
    done
}

get_encryption_choice() {
    step "ZFS Encryption"
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "ZFS native encryption protects your data at rest."
    echo ""
    echo "  - Passphrase required at every boot"
    echo "  - If forgotten, data is UNRECOVERABLE"
    echo "  - Recommended for laptops and sensitive data"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    prompt "Enable ZFS encryption? [Y/n]:"
    read -p "> " encrypt_choice

    if [[ "$encrypt_choice" =~ ^[Nn]$ ]]; then
        NO_ENCRYPT="yes"
        warn "Encryption DISABLED - data will not be encrypted at rest"
    else
        NO_ENCRYPT="no"
        info "Encryption enabled - you'll set a passphrase next"
    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 after install (key auth, fail2ban)."
    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 after install (key auth, fail2ban)!"
    fi
}

show_summary() {
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Configuration Summary:"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  Filesystem:   $FILESYSTEM"
    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)"
    if [[ "$FILESYSTEM" == "zfs" ]]; then
        if [[ "$NO_ENCRYPT" == "yes" ]]; then
            echo "  ZFS Pool:     $POOL_NAME (NOT encrypted)"
        else
            echo "  ZFS Pool:     $POOL_NAME (encrypted)"
        fi
        echo "  Boot:         ZFSBootMenu on all disks (redundant)"
    else
        if [[ "$NO_ENCRYPT" == "yes" ]]; then
            echo "  Encryption:   None"
        else
            echo "  Encryption:   LUKS2"
        fi
        echo "  Boot:         GRUB + grub-btrfs (snapshot boot)"
    fi
    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: 512M EFI + rest for ZFS
        # EFI only needs to hold ZFSBootMenu binary (~64MB) - 512MB is plenty
        sgdisk -n 1:0:+512M -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 pool (with or without encryption)
    # Note: We use zfs-import-scan at boot which doesn't require a cachefile
    if [[ "$NO_ENCRYPT" == "yes" ]]; then
        warn "Creating pool WITHOUT encryption (testing mode)"
        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 mountpoint=none \
            -R /mnt \
            "$POOL_NAME" $pool_config
    else
        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
    fi

    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"
    # EFI partition mounts at /efi - only holds ZFSBootMenu binary
    # /boot is a directory on ZFS root (kernels live on ZFS for snapshot safety)
    mkdir -p /mnt/efi
    mount "${EFI_PARTS[0]}" /mnt/efi
    info "EFI partition ${EFI_PARTS[0]} mounted at /mnt/efi"
}

install_base() {
    step "Installing Base System"

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

    # Add archzfs repo to pacman.conf for pacstrap
    # SigLevel=Never: pacstrap -K creates empty keyring where key import fails;
    # repo is explicitly added and served over HTTPS, GPG adds no real value here
    if ! grep -q "\[archzfs\]" /etc/pacman.conf; then
        cat >> /etc/pacman.conf << 'EOF'

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

    info "Installing base packages (this takes a while)..."
    info "ZFS will be built from source via DKMS - this ensures kernel compatibility."
    # Use yes to auto-select defaults for provider prompts
    yes "" | pacstrap -K /mnt \
        base \
        base-devel \
        linux-lts \
        linux-lts-headers \
        linux-firmware \
        zfs-dkms \
        zfs-utils \
        efibootmgr \
        networkmanager \
        avahi \
        nss-mdns \
        openssh \
        git \
        vim \
        sudo \
        zsh \
        nodejs \
        npm \
        ttf-dejavu \
        fzf \
        wget \
        wireless-regdb

    info "Base system installed."
}

install_base_btrfs() {
    step "Installing Base System (Btrfs)"

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

    info "Installing base packages (this takes a while)..."
    yes "" | pacstrap -K /mnt \
        base \
        base-devel \
        linux-lts \
        linux-lts-headers \
        linux-firmware \
        btrfs-progs \
        grub \
        grub-btrfs \
        efibootmgr \
        snapper \
        snap-pac \
        networkmanager \
        avahi \
        nss-mdns \
        openssh \
        git \
        vim \
        sudo \
        zsh \
        nodejs \
        npm \
        ttf-dejavu \
        fzf \
        wget \
        wireless-regdb

    info "Base system installed."
}

configure_system() {
    step "Configuring System"

    # fstab (only for EFI - /boot is on ZFS root)
    info "Generating fstab..."
    echo "# /efi - EFI System Partition (ZFSBootMenu binary)" > /mnt/etc/fstab
    echo "UUID=$(blkid -s UUID -o value "${EFI_PARTS[0]}")  /efi  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 (SigLevel=Never — same rationale as install_base)
    info "Adding archzfs repository..."
    cat >> /mnt/etc/pacman.conf << 'EOF'

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

    # Configure journald for ZFS
    # Problem: journald starts before ZFS mounts /var/log, so journal files
    # get created in tmpfs then hidden when ZFS mounts over it.
    # Solution: Make journal-flush wait for zfs-mount, and enable persistent storage.
    info "Configuring journald for ZFS..."
    mkdir -p /mnt/etc/systemd/journald.conf.d
    cat > /mnt/etc/systemd/journald.conf.d/persistent.conf << 'EOF'
[Journal]
Storage=persistent
EOF

    mkdir -p /mnt/etc/systemd/system/systemd-journal-flush.service.d
    cat > /mnt/etc/systemd/system/systemd-journal-flush.service.d/zfs.conf << 'EOF'
[Unit]
After=zfs-mount.service
EOF

    # 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 "Harden SSH after install (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

    # CRITICAL: Remove archiso drop-in that overrides mkinitcpio.conf HOOKS
    # The archiso.conf contains live ISO-specific hooks that are incompatible with ZFS
    # If not removed, it overrides our HOOKS setting and breaks boot after kernel updates
    if [[ -f /mnt/etc/mkinitcpio.conf.d/archiso.conf ]]; then
        info "Removing archiso drop-in config..."
        rm -f /mnt/etc/mkinitcpio.conf.d/archiso.conf
    fi

    # CRITICAL: Fix linux-lts preset file
    # The preset from archiso uses archiso-specific config that breaks mkinitcpio -P
    info "Creating proper linux-lts preset..."
    cat > /mnt/etc/mkinitcpio.d/linux-lts.preset << 'PRESET_EOF'
# mkinitcpio preset file for linux-lts

PRESETS=(default fallback)

ALL_kver="/boot/vmlinuz-linux-lts"

default_image="/boot/initramfs-linux-lts.img"

fallback_image="/boot/initramfs-linux-lts-fallback.img"
fallback_options="-S autodetect"
PRESET_EOF

    # Check for AMD ISP (Image Signal Processor) firmware needs
    # ISP is used for camera processing on AMD APUs (Strix, Strix Halo, etc.)
    # The firmware must be in initramfs since amdgpu loads before root is mounted
    if lspci | grep -qi "amd.*display\|amd.*vga\|radeon"; then
        local isp_firmware
        isp_firmware=$(ls /mnt/usr/lib/firmware/amdgpu/isp_*.bin.zst 2>/dev/null | head -1)
        if [[ -n "$isp_firmware" ]]; then
            # Remove /mnt prefix - config is used inside chroot where root is /
            local chroot_path="${isp_firmware#/mnt}"
            info "AMD APU detected with ISP firmware - adding to initramfs"
            mkdir -p /mnt/etc/mkinitcpio.conf.d
            cat > /mnt/etc/mkinitcpio.conf.d/amd-isp.conf << EOF
# AMD ISP (Image Signal Processor) firmware for camera support
# Loaded early so amdgpu can initialize ISP before root is mounted
FILES+=($chroot_path)
EOF
        fi
    fi

    # Configure hooks for ZFS
    # - Use udev (not systemd): ZFS hook is busybox-based and incompatible with systemd init
    # - Remove autodetect: it filters modules based on live ISO hardware, not target
    #   This ensures NVMe, AHCI, and other storage drivers are always included
    # - Remove fsck: ZFS doesn't use it, avoids confusing error messages
    # - Add zfs: required for ZFS root boot
    sed -i 's/^HOOKS=.*/HOOKS=(base udev microcode modconf kms keyboard keymap consolefont block zfs filesystems)/' /mnt/etc/mkinitcpio.conf

    # Get the installed kernel version (not the running kernel)
    local kernel_ver
    kernel_ver=$(ls /mnt/usr/lib/modules | grep lts | head -1)
    if [[ -z "$kernel_ver" ]]; then
        error "Could not find LTS kernel modules"
    fi
    info "Installed kernel: $kernel_ver"

    # Ensure kernel module dependencies are up to date after DKMS build
    # Must specify kernel version since running kernel differs from installed kernel
    info "Updating module dependencies..."
    arch-chroot /mnt depmod "$kernel_ver"

    # Verify ZFS module exists
    if ! [[ -f "/mnt/usr/lib/modules/$kernel_ver/updates/dkms/zfs.ko.zst" ]]; then
        error "ZFS module not found! DKMS build may have failed."
    fi
    info "ZFS module verified for kernel $kernel_ver"

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

configure_zfsbootmenu() {
    step "Configuring ZFSBootMenu"

    # Ensure hostid exists BEFORE reading it
    # This is critical: hostid command returns a value even without /etc/hostid,
    # but zgenhostid creates a DIFFERENT value. We must generate first, then read.
    if [[ ! -f /etc/hostid ]]; then
        zgenhostid
    fi

    # Now get the consistent hostid for kernel parameter
    local host_id
    host_id=$(hostid)

    # Copy hostid to installed system (ZFS uses this for pool ownership)
    cp /etc/hostid /mnt/etc/hostid

    # Create ZFSBootMenu directory on EFI
    mkdir -p /mnt/efi/EFI/ZBM

    # Download ZFSBootMenu release EFI binary
    # Using the bundled release which includes everything needed
    info "Downloading ZFSBootMenu..."
    local zbm_url="https://get.zfsbootmenu.org/efi"
    if ! curl -fsSL -o /mnt/efi/EFI/ZBM/zfsbootmenu.efi "$zbm_url"; then
        error "Failed to download ZFSBootMenu"
    fi
    info "ZFSBootMenu binary installed."

    # Set kernel command line on the ROOT PARENT dataset
    # This allows inheritance to all boot environments (future-proofing)
    # ZFSBootMenu reads org.zfsbootmenu:commandline property
    local cmdline="rw loglevel=3"

    # Add any AMD GPU workarounds if needed (detect Strix Halo etc)
    if lspci | grep -qi "amd.*display\|amd.*vga"; then
        info "AMD GPU detected - adding workaround parameters"
        cmdline="$cmdline amdgpu.pg_mask=0 amdgpu.cwsr_enable=0"
    fi

    # Set on ROOT parent so all boot environments inherit it
    zfs set org.zfsbootmenu:commandline="$cmdline" "$POOL_NAME/ROOT"
    info "Kernel command line set on $POOL_NAME/ROOT (inherited by children)"

    # Set bootfs property - tells ZFSBootMenu which dataset to boot by default
    zpool set bootfs="$POOL_NAME/ROOT/default" "$POOL_NAME"
    info "Default boot filesystem set to $POOL_NAME/ROOT/default"

    # Create EFI boot entries for each disk
    # ZFSBootMenu EFI parameters (passed via --unicode):
    #   spl_hostid=0x...     - Required for pool import
    #   zbm.timeout=3        - Seconds before auto-boot (-1 = always show menu)
    #   zbm.prefer=POOLNAME  - Preferred pool to boot from
    #   zbm.import_policy=hostid - How to handle pool imports
    local zbm_cmdline="spl_hostid=0x${host_id} zbm.timeout=3 zbm.prefer=${POOL_NAME} zbm.import_policy=hostid"

    for i in "${!SELECTED_DISKS[@]}"; do
        local disk="${SELECTED_DISKS[$i]}"
        local label="ZFSBootMenu"
        if [[ ${#SELECTED_DISKS[@]} -gt 1 ]]; then
            label="ZFSBootMenu-disk$((i+1))"
        fi

        # Determine partition number (always 1 - first partition is EFI)
        local part_num=1

        info "Creating EFI boot entry: $label on $disk"
        efibootmgr --create \
            --disk "$disk" \
            --part "$part_num" \
            --label "$label" \
            --loader '\EFI\ZBM\zfsbootmenu.efi' \
            --unicode "$zbm_cmdline" \
            --quiet
    done

    # Get the boot entry number and set as first in boot order
    local bootnum
    bootnum=$(efibootmgr | grep "ZFSBootMenu" | head -1 | grep -oP 'Boot\K[0-9A-F]+')
    if [[ -n "$bootnum" ]]; then
        # Get current boot order, prepend our entry
        local current_order
        current_order=$(efibootmgr | grep "BootOrder" | cut -d: -f2 | tr -d ' ')
        efibootmgr --bootorder "$bootnum,$current_order" --quiet
        info "ZFSBootMenu set as primary boot option"
    fi

    info "ZFSBootMenu configuration complete."
}

configure_zfs_services() {
    step "Configuring ZFS Services"

    arch-chroot /mnt systemctl enable zfs.target

    # Use zfs-import-scan instead of zfs-import-cache
    # This is the recommended method - it uses blkid to scan for pools
    # and doesn't require a cachefile
    # Note: ZFS package preset enables zfs-import-cache by default, so we must
    # explicitly disable it before enabling zfs-import-scan
    arch-chroot /mnt systemctl disable zfs-import-cache.service
    arch-chroot /mnt systemctl enable zfs-import-scan.service
    arch-chroot /mnt systemctl enable zfs-mount.service
    arch-chroot /mnt systemctl enable zfs-import.target

    # Note: hostid and bootfs are already set by configure_zfsbootmenu()

    # Disable cachefile - we use zfs-import-scan which doesn't need it
    # Also remove any existing cachefile since zfs-import-scan has a condition
    # that prevents it from running if /etc/zfs/zpool.cache exists
    zpool set cachefile=none "$POOL_NAME"
    rm -f /mnt/etc/zfs/zpool.cache

    # Enable other services
    arch-chroot /mnt systemctl enable NetworkManager
    arch-chroot /mnt systemctl enable avahi-daemon
    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."
}

configure_zfs_tools() {
    step "Installing ZFS Management Tools"

    # Copy ZFS management scripts
    cp /usr/local/bin/zfssnapshot /mnt/usr/local/bin/zfssnapshot
    cp /usr/local/bin/zfsrollback /mnt/usr/local/bin/zfsrollback
    chmod +x /mnt/usr/local/bin/zfssnapshot
    chmod +x /mnt/usr/local/bin/zfsrollback

    info "ZFS management scripts installed: zfssnapshot, zfsrollback"
    info "Tip: Install sanoid for automated snapshot retention."
}

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

    step "Syncing EFI Partitions for Redundancy"

    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 ZFSBootMenu to EFI partition $((i+1)): $efi_part"

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

        # Copy ZFSBootMenu binary to secondary EFI partitions
        mkdir -p "$temp_mount/EFI/ZBM"
        cp /mnt/efi/EFI/ZBM/zfsbootmenu.efi "$temp_mount/EFI/ZBM/"

        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/efi 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"
    if [[ "$NO_ENCRYPT" == "yes" ]]; then
        echo "  ZFS Pool:     $POOL_NAME (not encrypted)"
    else
        echo "  ZFS Pool:     $POOL_NAME (encrypted)"
    fi
    echo ""
    echo "ZFSBootMenu Features:"
    echo "  - Boot from any snapshot (Ctrl+D at boot menu)"
    echo "  - Genesis snapshot: pristine post-install state"
    echo "  - Pre-pacman snapshots for safe upgrades"
    echo "  - Install sanoid/syncoid for automated retention"
    echo ""
    echo "Boot Menu Keys (at ZFSBootMenu):"
    echo "  Enter    - Boot selected environment"
    echo "  e        - Edit kernel command line"
    echo "  Ctrl+D   - Show snapshot selector"
    echo "  Ctrl+R   - Recovery shell"
    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 ""
}

print_btrfs_summary() {
    echo ""
    echo "╔═══════════════════════════════════════════════════════════════╗"
    echo "║                    Installation Complete!                     ║"
    echo "╚═══════════════════════════════════════════════════════════════╝"
    echo ""
    echo "System Configuration:"
    echo "  Hostname:     $HOSTNAME"
    echo "  Timezone:     $TIMEZONE"
    echo "  Filesystem:   Btrfs"
    if [[ "$NO_ENCRYPT" == "yes" ]]; then
        echo "  Encryption:   None"
    else
        echo "  Encryption:   LUKS2"
    fi
    echo ""
    echo "Btrfs Snapshot Features:"
    echo "  - Boot from any snapshot via GRUB menu"
    echo "  - Genesis snapshot: pristine post-install state"
    echo "  - Pre/post pacman snapshots via snap-pac"
    echo "  - Timeline snapshots: 6 hourly, 7 daily, 2 weekly, 1 monthly"
    echo ""
    echo "GRUB Boot Menu:"
    echo "  - Select 'Arch Linux snapshots' submenu to boot from snapshots"
    echo "  - Snapshots auto-added when created by snapper"
    echo ""
    echo "Useful Commands:"
    echo "  List snapshots:    snapper -c root list"
    echo "  Manual snapshot:   snapper -c root create -d 'description'"
    echo "  Rollback:          snapper -c root rollback <number>"
    echo "  Compare:           snapper -c root diff <num1>..<num2>"
    echo ""
    info "Installation log: $LOGFILE"
    echo ""
}

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

main() {
    parse_args "$@"
    preflight_checks
    check_config
    gather_input
    filesystem_preflight

    # Unattended installation begins
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Beginning unattended installation ($FILESYSTEM)..."
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    if [[ "$FILESYSTEM" == "zfs" ]]; then
        install_zfs
    elif [[ "$FILESYSTEM" == "btrfs" ]]; then
        install_btrfs
    fi
}

#############################
# ZFS Installation Path
#############################

install_zfs() {
    partition_disks
    create_zfs_pool
    create_datasets
    mount_efi
    install_base
    configure_system
    configure_wifi
    configure_ssh
    configure_initramfs
    configure_zfsbootmenu
    configure_zfs_services
    configure_pacman_hook
    configure_zfs_tools
    sync_efi_partitions
    create_genesis_snapshot
    cleanup
    print_summary
}

#############################
# Btrfs Installation Path
#############################

install_btrfs() {
    local num_disks=${#SELECTED_DISKS[@]}
    local btrfs_devices=()
    local efi_parts=()
    local root_parts=()

    # Collect partition references for all disks
    for disk in "${SELECTED_DISKS[@]}"; do
        root_parts+=("$(get_root_partition "$disk")")
        efi_parts+=("$(get_efi_partition "$disk")")
    done

    # Partition all disks
    for disk in "${SELECTED_DISKS[@]}"; do
        partition_disk "$disk"
    done

    # Format all EFI partitions
    format_efi_partitions "${SELECTED_DISKS[@]}"

    # LUKS encryption (if enabled)
    if [[ "$NO_ENCRYPT" != "yes" ]]; then
        if [[ $num_disks -eq 1 ]]; then
            # Single disk LUKS
            create_luks_container "${root_parts[0]}" "$LUKS_PASSPHRASE"
            open_luks_container "${root_parts[0]}" "$LUKS_PASSPHRASE"
            btrfs_devices=("/dev/mapper/$LUKS_MAPPER_NAME")
        else
            # Multi-disk LUKS - encrypt each partition
            create_luks_containers "$LUKS_PASSPHRASE" "${root_parts[@]}"
            open_luks_containers "$LUKS_PASSPHRASE" "${root_parts[@]}"
            btrfs_devices=($(get_luks_devices $num_disks))
        fi
    else
        # No encryption - use raw partitions
        btrfs_devices=("${root_parts[@]}")
    fi

    # Create btrfs filesystem
    if [[ $num_disks -eq 1 ]]; then
        create_btrfs_volume "${btrfs_devices[0]}"
    else
        create_btrfs_volume "${btrfs_devices[@]}" --raid-level "$RAID_LEVEL"
    fi

    # Create and mount subvolumes (use first device for mount)
    create_btrfs_subvolumes "${btrfs_devices[0]}"
    mount_btrfs_subvolumes "${btrfs_devices[0]}"

    # Mount primary EFI
    mkdir -p /mnt/efi
    mount "${efi_parts[0]}" /mnt/efi

    # Install base system
    install_base_btrfs

    # Configure system
    configure_system
    configure_wifi
    configure_ssh

    # Configure encryption if enabled
    if [[ "$NO_ENCRYPT" != "yes" ]]; then
        setup_luks_testing_keyfile "$LUKS_PASSPHRASE" "${root_parts[@]}"
        configure_crypttab "${root_parts[@]}"
        configure_luks_grub "${root_parts[0]}"
        configure_luks_initramfs
    fi

    generate_btrfs_fstab "${btrfs_devices[0]}" "${efi_parts[0]}"
    configure_btrfs_initramfs

    # GRUB installation
    if [[ $num_disks -eq 1 ]]; then
        configure_grub "${efi_parts[0]}"
    else
        # Multi-disk: install GRUB to all EFI partitions
        configure_grub "${efi_parts[0]}"
        install_grub_all_efi "${efi_parts[@]}"
        create_grub_sync_hook "${efi_parts[@]}"
    fi

    configure_snapper
    configure_btrfs_services
    configure_btrfs_pacman_hook

    # Genesis snapshot
    create_btrfs_genesis_snapshot

    # Cleanup
    btrfs_cleanup
    if [[ "$NO_ENCRYPT" != "yes" ]]; then
        if [[ $num_disks -eq 1 ]]; then
            close_luks_container
        else
            close_luks_containers $num_disks
        fi
    fi
    print_btrfs_summary
}

trap 'error "Installation interrupted!"' INT TERM

main "$@"