#!/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 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 "  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 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."
    # 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
    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

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

    # 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 "Note: Install sanoid via archsetup for automated snapshot retention."
}

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 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 "  - Sanoid/syncoid configured by archsetup"
    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
    copy_archsetup
    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
    copy_archsetup

    # 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 "$@"