path: root/custom/install-archzfs

#!/bin/bash
# install-archzfs - Arch Linux ZFS Root Installation Script
# Craig Jennings <c@cjennings.net>
#
# Installs Arch Linux on ZFS root with native encryption.
# Designed to be run from the custom archzfs ISO.
#
# Features:
# - All questions asked upfront, then unattended installation
# - Optional WiFi configuration with connection test
# - ZFS native encryption (passphrase required at boot)
# - Pre-pacman ZFS snapshots for safe upgrades
#
# UNATTENDED MODE:
#   Use --config-file /path/to/install-archzfs.conf for automated installs.
#   Config file must be explicitly specified to prevent accidental disk wipes.
#   See /root/install-archzfs.conf.example for a template with all options.

set -e

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

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

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

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

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

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

#############################
# Config File Support
#############################

CONFIG_FILE=""
UNATTENDED=false

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --config-file)
                if [[ -n "$2" && ! "$2" =~ ^- ]]; then
                    CONFIG_FILE="$2"
                    shift 2
                else
                    error "--config-file requires a path argument"
                fi
                ;;
            --help|-h)
                echo "Usage: install-archzfs [OPTIONS]"
                echo ""
                echo "Options:"
                echo "  --config-file PATH  Use config file for unattended installation"
                echo "  --help, -h          Show this help message"
                echo ""
                echo "Without --config-file, runs in interactive mode."
                echo "See /root/install-archzfs.conf.example for a config template."
                exit 0
                ;;
            *)
                error "Unknown option: $1 (use --help for usage)"
                ;;
        esac
    done
}

load_config() {
    local config_path="$1"

    if [[ ! -f "$config_path" ]]; then
        error "Config file not found: $config_path"
    fi

    info "Loading config from: $config_path"

    # Source the config file (it's just key=value pairs)
    # shellcheck disable=SC1090
    source "$config_path"

    # Convert DISKS from comma-separated string to array
    if [[ -n "$DISKS" ]]; then
        IFS=',' read -ra SELECTED_DISKS <<< "$DISKS"
    fi

    UNATTENDED=true
    info "Running in unattended mode"
}

check_config() {
    # Only use config when explicitly specified with --config-file
    # This prevents accidental disk destruction from an unnoticed config file
    if [[ -n "$CONFIG_FILE" ]]; then
        load_config "$CONFIG_FILE"
    fi
}

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

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

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

    info "ZFS module loaded successfully."
}

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

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

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

        # 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 "  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"
        [[ -n "$WIFI_SSID" ]] && info "  WiFi: $WIFI_SSID"
        return 0
    fi

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

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

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

get_timezone() {
    step "Timezone"
    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_zfs_passphrase() {
    step "ZFS Encryption Passphrase"
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "This passphrase will be required at EVERY boot."
    echo ""
    echo "Requirements:"
    echo "  - Use a strong, memorable passphrase"
    echo "  - If forgotten, your data is UNRECOVERABLE"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

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

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

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

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

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

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

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

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

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

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

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

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

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

    EFI_PARTS=()
    ZFS_PARTS=()

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

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

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

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

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

        sleep 1
        partprobe "$disk"
    done

    sleep 2

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

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

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

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

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

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

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

create_datasets() {
    step "Creating ZFS Datasets"

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

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

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

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

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

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

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

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

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

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

install_base() {
    step "Installing Base System"

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

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

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

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

    info "Installing base packages (this takes a while)..."
    info "ZFS will be built from source via DKMS - this ensures kernel compatibility."
    # 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 \
        grub \
        freetype2 \
        efibootmgr \
        networkmanager \
        openssh \
        git \
        vim \
        sudo \
        zsh \
        nodejs \
        npm \
        ttf-dejavu

    info "Base system installed."
}

configure_system() {
    step "Configuring System"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

configure_initramfs() {
    step "Configuring Initramfs for ZFS"

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

    # Configure hooks for ZFS
    # ZFS doesn't use fsck - remove it to avoid confusing error messages
    sed -i 's/^HOOKS=.*/HOOKS=(base udev autodetect 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_bootloader() {
    step "Configuring GRUB Bootloader"

    # Get hostid for kernel parameter
    local host_id
    host_id=$(hostid)

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

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

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

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

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

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

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

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

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

configure_grub_zfs_snap() {
    step "Configuring ZFS Snapshot Boot Entries"

    # Install grub-zfs-snap script
    info "Installing grub-zfs-snap..."
    cp /usr/local/bin/grub-zfs-snap /mnt/usr/local/bin/grub-zfs-snap
    chmod +x /mnt/usr/local/bin/grub-zfs-snap

    # Install GRUB generator
    cp /usr/local/share/grub-zfs-snap/40_zfs_snapshots /mnt/etc/grub.d/40_zfs_snapshots
    chmod +x /mnt/etc/grub.d/40_zfs_snapshots

    # Install pacman hook for auto-regeneration
    mkdir -p /mnt/etc/pacman.d/hooks
    cp /usr/local/share/grub-zfs-snap/zz-grub-zfs-snap.hook /mnt/etc/pacman.d/hooks/

    info "ZFS snapshots will appear in GRUB boot menu."
    info "Run 'grub-zfs-snap' to manually regenerate after creating snapshots."
}

configure_zfs_services() {
    step "Configuring ZFS Services"

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

    # Copy hostid to installed system (ZFS uses this for pool ownership)
    if [[ -f /etc/hostid ]]; then
        cp /etc/hostid /mnt/etc/hostid
    else
        # Generate hostid if it doesn't exist
        zgenhostid
        cp /etc/hostid /mnt/etc/hostid
    fi

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

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

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

    info "ZFS services configured."
}

configure_pacman_hook() {
    step "Configuring Pacman Snapshot Hook"

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

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

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

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

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

# Prune old snapshots (runs quietly, non-blocking)
if [[ -x /usr/local/bin/zfs-snap-prune ]]; then
    /usr/local/bin/zfs-snap-prune --quiet &
fi
EOF

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

    info "Pacman hook configured."
}

configure_snapshot_retention() {
    step "Configuring Snapshot Retention"

    # Copy the prune script
    cp /usr/local/bin/zfs-snap-prune /mnt/usr/local/bin/zfs-snap-prune
    chmod +x /mnt/usr/local/bin/zfs-snap-prune

    # Create systemd service for pruning
    cat > /mnt/etc/systemd/system/zfs-snap-prune.service << 'EOF'
[Unit]
Description=Prune old ZFS snapshots
After=zfs.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/zfs-snap-prune --quiet
EOF

    # Create systemd timer for daily pruning
    cat > /mnt/etc/systemd/system/zfs-snap-prune.timer << 'EOF'
[Unit]
Description=Daily ZFS snapshot pruning

[Timer]
OnCalendar=daily
Persistent=true
RandomizedDelaySec=1h

[Install]
WantedBy=timers.target
EOF

    # Enable the timer
    arch-chroot /mnt systemctl enable zfs-snap-prune.timer

    info "Snapshot retention configured."
    info "Policy: Keep 20 recent, delete if older than 180 days"
    info "Genesis snapshot is always preserved."
}

copy_archsetup() {
    step "Installing archsetup Launcher"

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

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

    step "Syncing EFI Partitions for Redundancy"

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

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

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

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

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

        umount "$temp_mount"
    done

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

create_genesis_snapshot() {
    step "Creating Genesis Snapshot"

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

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

set -e

POOL_NAME="zroot"

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

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

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

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

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

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

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

cleanup() {
    step "Cleaning Up"

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

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

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

    info "Cleanup complete."
}

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

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

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

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

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

trap 'error "Installation interrupted!"' INT TERM

main "$@"