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

set -e

# These will be set interactively
HOSTNAME=""
USERNAME=""
TIMEZONE=""
LOCALE="en_US.UTF-8"
KEYMAP="us"

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

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

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

info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }
step()  { echo -e "\n${BLUE}==>${NC} ${CYAN}$1${NC}"; }
prompt() { echo -e "${BOLD}$1${NC}"; }

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

### Interactive Configuration ###
configure_install() {
    step "Installation Configuration"
    echo ""

    # 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

    echo ""

    # Username
    prompt "Enter primary username:"
    read -p "> " USERNAME
    while [[ -z "$USERNAME" || ! "$USERNAME" =~ ^[a-z_][a-z0-9_-]*$ ]]; do
        warn "Invalid username. Use lowercase letters, numbers, underscore, hyphen."
        read -p "> " USERNAME
    done

    echo ""

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

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

    echo ""

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

    echo ""

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

    # Confirm settings
    echo ""
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}Configuration Summary:${NC}"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo "  Hostname:  $HOSTNAME"
    echo "  Username:  $USERNAME"
    echo "  Timezone:  $TIMEZONE"
    echo "  Locale:    $LOCALE"
    echo "  Keymap:    $KEYMAP"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""

    read -p "Is this correct? [Y/n]: " confirm
    if [[ "$confirm" == "n" || "$confirm" == "N" ]]; then
        configure_install
    fi
}

### Disk Selection ###
select_disk() {
    step "Disk Selection"

    echo ""
    echo "Available disks:"
    echo "----------------"
    lsblk -d -o NAME,SIZE,MODEL,TYPE | grep disk
    echo ""

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

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

    PS3="Select disk for installation (number): "
    select disk in "${DISKS[@]}"; do
        if [[ -n "$disk" ]]; then
            DISK="/dev/$disk"
            break
        fi
    done

    echo ""
    warn "Selected disk: $DISK"
    echo ""
    lsblk "$DISK"
    echo ""

    read -p "This will DESTROY all data on $DISK. Type 'yes' to continue: " confirm
    [[ "$confirm" != "yes" ]] && error "Aborted by user"
}

### Partitioning ###
partition_disk() {
    step "Partitioning $DISK"

    # Wipe existing signatures
    info "Wiping existing signatures..."
    wipefs -af "$DISK"
    sgdisk --zap-all "$DISK"

    # Create partitions
    # 1: EFI System Partition (1GB)
    # 2: ZFS partition (rest)
    info "Creating partitions..."
    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 vs sda naming)
    if [[ "$DISK" == *"nvme"* ]] || [[ "$DISK" == *"mmcblk"* ]]; then
        EFI_PART="${DISK}p1"
        ZFS_PART="${DISK}p2"
    else
        EFI_PART="${DISK}1"
        ZFS_PART="${DISK}2"
    fi

    # Wait for partitions to appear
    sleep 2
    partprobe "$DISK"
    sleep 2

    # Format EFI partition
    info "Formatting EFI partition..."
    mkfs.fat -F32 -n EFI "$EFI_PART"

    info "Partitioning complete."
    lsblk "$DISK"
}

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

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

    echo ""
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}ZFS Encryption Passphrase${NC}"
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""
    echo "You will now create an encryption passphrase."
    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 -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""

    # Create encrypted pool
    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" "$ZFS_PART"

    info "ZFS pool created successfully."
}

### Dataset Creation ###
create_datasets() {
    step "Creating ZFS Datasets"

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

    # Main root filesystem with reservation for safety
    zfs create -o mountpoint=/ -o canmount=noauto -o reservation=50G "$POOL_NAME/ROOT/default"

    # Mount root first
    zfs mount "$POOL_NAME/ROOT/default"

    # Home datasets
    zfs create -o mountpoint=/home "$POOL_NAME/home"
    zfs create -o mountpoint=/root "$POOL_NAME/home/root"
    zfs create -o mountpoint="/home/$USERNAME" "$POOL_NAME/home/$USERNAME"

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

    # VMs dataset - larger 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"

    # Exclude temp directories 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:"
    echo ""
    zfs list -r "$POOL_NAME" -o name,mountpoint,compression,reservation
}

### Mount EFI ###
mount_efi() {
    step "Mounting EFI Partition"

    mkdir -p /mnt/boot
    mount "$EFI_PART" /mnt/boot

    info "EFI partition mounted at /mnt/boot"
}

### Install Base System ###
install_base() {
    step "Installing Base System"

    info "Updating pacman keys..."
    pacman-key --init
    pacman-key --populate archlinux
    pacman-key -r DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true
    pacman-key --lsign-key DDF7DB817396A49B2A2723F7403BD972F75D9D76 2>/dev/null || true

    info "Installing base packages (this takes a while)..."
    pacstrap -K /mnt \
        base \
        base-devel \
        linux \
        linux-headers \
        linux-firmware \
        zfs-linux \
        zfs-utils \
        grub \
        efibootmgr \
        networkmanager \
        openssh \
        git \
        vim \
        sudo \
        zsh \
        nodejs \
        npm \
        sanoid

    info "Base system installed."
}

### Configure System ###
configure_system() {
    step "Configuring System"

    # Generate fstab (only for EFI, ZFS handles the rest)
    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 to installed system
    info "Adding archzfs repository..."
    cat >> /mnt/etc/pacman.conf << 'EOF'

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

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

### Configure mkinitcpio ###
configure_initramfs() {
    step "Configuring Initramfs for ZFS"

    # Backup original
    cp /mnt/etc/mkinitcpio.conf /mnt/etc/mkinitcpio.conf.bak

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

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

### Configure Bootloader ###
configure_bootloader() {
    step "Configuring GRUB Bootloader"

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

    info "Installing GRUB..."
    arch-chroot /mnt grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB

    info "Generating GRUB configuration..."
    arch-chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
}

### Configure ZFS Services ###
configure_zfs_services() {
    step "Configuring ZFS Services"

    # Enable ZFS services
    arch-chroot /mnt systemctl enable zfs.target
    arch-chroot /mnt systemctl enable zfs-import-cache
    arch-chroot /mnt systemctl enable zfs-mount
    arch-chroot /mnt systemctl enable zfs-import.target

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

    # Set bootfs property
    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 Sanoid (Snapshot Management) ###
configure_sanoid() {
    step "Configuring Sanoid Snapshot Management"

    mkdir -p /mnt/etc/sanoid

    cat > /mnt/etc/sanoid/sanoid.conf << EOF
# Sanoid configuration for ZFS snapshots
# https://github.com/jimsalterjrs/sanoid

#############################
# Templates
#############################

[template_production]
    # Frequent snapshots for active data
    hourly = 24
    daily = 7
    weekly = 4
    monthly = 12
    yearly = 0
    autosnap = yes
    autoprune = yes

[template_backup]
    # Less frequent for large/static data
    hourly = 0
    daily = 7
    weekly = 4
    monthly = 6
    yearly = 0
    autosnap = yes
    autoprune = yes

[template_none]
    # No automatic snapshots (for tmp, cache)
    autosnap = no
    autoprune = yes

#############################
# Datasets
#############################

# Root filesystem
[$POOL_NAME/ROOT/default]
    use_template = production

# Home directories
[$POOL_NAME/home]
    use_template = production
    recursive = yes

# Media (large files, less frequent snapshots)
[$POOL_NAME/media]
    use_template = backup

# VMs (snapshot before changes manually, or less frequently)
[$POOL_NAME/vms]
    use_template = backup

# Var data
[$POOL_NAME/var/log]
    use_template = production

[$POOL_NAME/var/lib/pacman]
    use_template = production

# No snapshots for cache/tmp (handled by dataset property, but explicit here)
[$POOL_NAME/var/cache]
    use_template = none

[$POOL_NAME/var/tmp]
    use_template = none

[$POOL_NAME/tmp]
    use_template = none
EOF

    # Enable sanoid timer
    arch-chroot /mnt systemctl enable sanoid.timer

    info "Sanoid configured. Snapshots will run automatically."
}

### Configure Pacman ZFS Snapshot Hook ###
configure_pacman_hook() {
    step "Configuring Pacman Pre-Upgrade 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
# Create a ZFS snapshot before pacman transactions
# This allows easy rollback if an upgrade breaks something

POOL="zroot"
DATASET="$POOL/ROOT/default"
TIMESTAMP=$(date +%Y-%m-%d_%H-%M-%S)
SNAPSHOT_NAME="pre-pacman_$TIMESTAMP"

# Create the snapshot
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. Snapshots will be created before each transaction."
}

### Create User ###
create_user() {
    step "Creating User: $USERNAME"

    arch-chroot /mnt useradd -m -G wheel -s /bin/zsh "$USERNAME" 2>/dev/null || \
        warn "User $USERNAME may already exist"

    # Set ownership of home dataset
    arch-chroot /mnt chown -R "$USERNAME:$USERNAME" "/home/$USERNAME"

    # Configure sudo
    echo "%wheel ALL=(ALL:ALL) ALL" > /mnt/etc/sudoers.d/wheel
    chmod 440 /mnt/etc/sudoers.d/wheel

    echo ""
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}Set User Password${NC}"
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    info "Set password for $USERNAME:"
    arch-chroot /mnt passwd "$USERNAME"

    echo ""
    info "Set password for root:"
    arch-chroot /mnt passwd
}

### Copy archsetup ###
copy_archsetup() {
    step "Copying archsetup to New System"

    if [[ -d /code/archsetup ]]; then
        mkdir -p "/mnt/home/$USERNAME/code"
        cp -r /code/archsetup "/mnt/home/$USERNAME/code/"
        arch-chroot /mnt chown -R "$USERNAME:$USERNAME" "/home/$USERNAME/code"
        info "archsetup copied to /home/$USERNAME/code/archsetup"
    else
        warn "archsetup not found in ISO, skipping..."
    fi
}

### Create Syncoid Script for TrueNAS ###
create_syncoid_script() {
    step "Creating Syncoid Replication Script"

    cat > /mnt/usr/local/bin/zfs-replicate << 'SCRIPT'
#!/bin/bash
# zfs-replicate - Replicate ZFS datasets to TrueNAS
# Usage: zfs-replicate [dataset] [target]
#
# Examples:
#   zfs-replicate                           # Replicate all configured datasets
#   zfs-replicate zroot/home user@truenas:/tank/backup/laptop

set -e

# Configuration - edit these for your TrueNAS setup
TRUENAS_HOST="truenas"           # TrueNAS hostname or IP
TRUENAS_USER="root"              # User with ZFS permissions
TRUENAS_POOL="tank"              # Destination pool
BACKUP_PATH="backup/laptop"      # Path under the pool

# Datasets to replicate (space-separated)
DATASETS="zroot/ROOT/default zroot/home zroot/media zroot/vms"

# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m'

info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

# Check if syncoid is installed
command -v syncoid >/dev/null 2>&1 || error "syncoid not found. Install sanoid package."

# Single dataset mode
if [[ -n "$1" ]] && [[ -n "$2" ]]; then
    info "Replicating $1 to $2"
    syncoid --recursive "$1" "$2"
    exit 0
fi

# Full replication mode
info "Starting ZFS replication to $TRUENAS_HOST"
echo ""

for dataset in $DATASETS; do
    dest="$TRUENAS_USER@$TRUENAS_HOST:$TRUENAS_POOL/$BACKUP_PATH/${dataset#zroot/}"
    info "Replicating $dataset -> $dest"

    if syncoid --recursive "$dataset" "$dest"; then
        info "  Success"
    else
        warn "  Failed (will retry next run)"
    fi
    echo ""
done

info "Replication complete."
SCRIPT

    chmod +x /mnt/usr/local/bin/zfs-replicate

    # Create systemd service and timer for automatic replication
    cat > /mnt/etc/systemd/system/zfs-replicate.service << 'EOF'
[Unit]
Description=ZFS Replication to TrueNAS
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/zfs-replicate
User=root

[Install]
WantedBy=multi-user.target
EOF

    cat > /mnt/etc/systemd/system/zfs-replicate.timer << 'EOF'
[Unit]
Description=Run ZFS replication nightly

[Timer]
OnCalendar=*-*-* 02:00:00
RandomizedDelaySec=1800
Persistent=true

[Install]
WantedBy=timers.target
EOF

    info "Syncoid replication script created."
    info "Edit /usr/local/bin/zfs-replicate to configure your TrueNAS connection."
    info "Enable with: systemctl enable --now zfs-replicate.timer"
}

### Unmount and Export ###
cleanup() {
    step "Cleaning Up"

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

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

    info "Cleanup complete."
}

### Print Summary ###
print_summary() {
    echo ""
    echo -e "${GREEN}╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}║                    Installation Complete!                     ║${NC}"
    echo -e "${GREEN}╚═══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "${BOLD}System Configuration:${NC}"
    echo "  Hostname:     $HOSTNAME"
    echo "  Username:     $USERNAME"
    echo "  Timezone:     $TIMEZONE"
    echo "  ZFS Pool:     $POOL_NAME (encrypted)"
    echo ""
    echo -e "${BOLD}ZFS Features Configured:${NC}"
    echo "  - Automatic snapshots via sanoid (hourly/daily/weekly/monthly)"
    echo "  - Pre-pacman snapshots for safe upgrades"
    echo "  - Replication script ready for TrueNAS"
    echo ""
    echo -e "${BOLD}Next Steps:${NC}"
    echo "  1. Reboot: ${CYAN}reboot${NC}"
    echo "  2. Enter your ZFS encryption passphrase at boot"
    echo "  3. Log in as $USERNAME"
    echo "  4. Run archsetup: ${CYAN}cd ~/code/archsetup && sudo ./archsetup${NC}"
    echo ""
    echo -e "${BOLD}Configure TrueNAS Replication:${NC}"
    echo "  1. Set up SSH key auth to TrueNAS"
    echo "  2. Edit: ${CYAN}/usr/local/bin/zfs-replicate${NC}"
    echo "  3. Enable: ${CYAN}sudo systemctl enable --now zfs-replicate.timer${NC}"
    echo ""
    echo -e "${BOLD}Useful ZFS Commands:${NC}"
    echo "  List snapshots:     ${CYAN}zfs list -t snapshot${NC}"
    echo "  Manual snapshot:    ${CYAN}sudo zfs snapshot zroot/home@my-snapshot${NC}"
    echo "  Rollback:           ${CYAN}sudo zfs rollback zroot/home@my-snapshot${NC}"
    echo "  Check pool status:  ${CYAN}zpool status${NC}"
    echo ""
    echo -e "${BOLD}If Something Goes Wrong:${NC}"
    echo "  Boot from this ISO, then:"
    echo "    ${CYAN}zpool import -R /mnt zroot${NC}"
    echo "    ${CYAN}zfs load-key zroot${NC}"
    echo "    ${CYAN}zfs mount zroot/ROOT/default${NC}"
    echo ""
    info "Installation log saved to: $LOGFILE"
    echo ""
}

### Main Installation Flow ###
main() {
    echo ""
    echo -e "${CYAN}╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║       Arch Linux ZFS Root Installation                        ║${NC}"
    echo -e "${CYAN}║       with Native Encryption                                  ║${NC}"
    echo -e "${CYAN}╚═══════════════════════════════════════════════════════════════╝${NC}"
    echo ""

    info "Installation log: $LOGFILE"
    echo ""

    configure_install
    select_disk
    partition_disk
    create_zfs_pool
    create_datasets
    mount_efi
    install_base
    configure_system
    configure_initramfs
    configure_bootloader
    configure_zfs_services
    configure_sanoid
    configure_pacman_hook
    create_user
    copy_archsetup
    create_syncoid_script
    cleanup
    print_summary
}

# Handle interrupts
trap 'error "Installation interrupted!"' INT TERM

# Run main
main "$@"