#!/bin/sh
# archsetup-zfs - Post-installation setup for Arch Linux on ZFS
# Craig Jennings <craigmartinjennings@gmail.com>
# License: GNU GPLv3
#
# This is a ZFS-specific variant of archsetup.
# It replaces btrfs snapshot tooling with ZFS equivalents.
#
# Run this after install-archzfs completes and you've rebooted.

# Commentary
#
# This script is based on the original archsetup but modified for ZFS:
# - Removes: timeshift-autosnap, grub-btrfs (btrfs-specific)
# - Adds: sanoid/syncoid configuration verification
# - Keeps: All other package installations and configurations
#
# The ZFS snapshot infrastructure is already configured by install-archzfs:
# - sanoid.timer for automatic snapshots
# - pacman hook for pre-upgrade snapshots
# - syncoid script for TrueNAS replication

set -e

### Root Check
if [ "$EUID" -ne 0 ]; then
    echo "ERROR: This script must be run as root"
    echo "Usage: sudo $0"
    exit 1
fi

### Parse Arguments
skip_slow_packages=false
fresh_install=false
show_status_only=false

while [ $# -gt 0 ]; do
    case "$1" in
        --skip-slow-packages)
            skip_slow_packages=true
            shift
            ;;
        --fresh)
            fresh_install=true
            shift
            ;;
        --status)
            show_status_only=true
            shift
            ;;
        --help|-h)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --skip-slow-packages  Skip texlive-meta and topgrade"
            echo "  --fresh               Start fresh, ignore previous progress"
            echo "  --status              Show installation progress and exit"
            echo "  --help, -h            Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            echo "Usage: $0 [--skip-slow-packages] [--fresh] [--status]"
            exit 1
            ;;
    esac
done

### Constants
username="cjennings"
password="welcome"  # will be changed on first login

archsetup_dir="$(dirname "$(readlink -f "$0")")"

# Check if we're running from the ISO copy or user's home
if [ -d "/home/$username/code/archsetup" ]; then
    archsetup_dir="/home/$username/code/archsetup"
fi

dotfiles_dir="$archsetup_dir/dotfiles"

dwm_repo="https://git.cjennings.net/dwm.git"
dmenu_repo="https://git.cjennings.net/dmenu.git"
st_repo="https://git.cjennings.net/st.git"
slock_repo="https://git.cjennings.net/slock.git"
dotemacs_repo="https://git.cjennings.net/dotemacs.git"

logfile="/var/log/archsetup-zfs-$(date +'%Y-%m-%d-%H-%M-%S').log"
source_dir="/home/$username/.local/src"
packages_before="/var/log/archsetup-preexisting-package-list.txt"
packages_after="/var/log/archsetup-post-install-package-list.txt"
archsetup_packages="/var/log/archsetup-installed-packages.txt"

min_disk_space_gb=20
state_dir="/var/lib/archsetup-zfs/state"

### ZFS Check
check_zfs() {
    if ! command -v zfs >/dev/null 2>&1; then
        echo "ERROR: ZFS tools not found. This script is for ZFS systems."
        echo "Use the regular 'archsetup' for btrfs systems."
        exit 1
    fi

    if ! zpool list zroot >/dev/null 2>&1; then
        echo "ERROR: ZFS pool 'zroot' not found."
        echo "This script expects a ZFS root pool named 'zroot'."
        exit 1
    fi

    echo "ZFS system detected: $(zpool list -H -o name,size,health zroot)"
}

### State Tracking
step_completed() {
    [ -f "$state_dir/$1" ]
}

mark_complete() {
    mkdir -p "$state_dir"
    echo "$(date +'%Y-%m-%d %H:%M:%S')" > "$state_dir/$1"
}

run_step() {
    step_name="$1"
    step_func="$2"

    if step_completed "$step_name"; then
        printf "Skipping %s (already completed)\n" "$step_name"
        return 0
    fi

    if $step_func; then
        mark_complete "$step_name"
        return 0
    else
        printf "FAILED: %s\n" "$step_name"
        printf "To retry this step, remove: %s/%s\n" "$state_dir" "$step_name"
        return 1
    fi
}

show_status() {
    echo "Archsetup-ZFS State Status"
    echo "=========================="
    echo "State directory: $state_dir"
    echo ""
    if [ ! -d "$state_dir" ]; then
        echo "No state found. Script has not been run or was run with --fresh."
        exit 0
    fi
    echo "Completed steps:"
    for step in intro prerequisites user_customizations \
                aur_installer essential_services xorg dwm \
                desktop_environment developer_workstation \
                supplemental_software boot_ux zfs_services; do
        if step_completed "$step"; then
            timestamp=$(cat "$state_dir/$step")
            printf "  [x] %-25s (%s)\n" "$step" "$timestamp"
        else
            printf "  [ ] %-25s\n" "$step"
        fi
    done
    exit 0
}

if $show_status_only; then
    show_status
fi

if $fresh_install; then
    echo "Starting fresh installation (removing previous state)..."
    rm -rf "$state_dir"
fi

### General Functions
error () {
    case "$1" in
        "error")
            printf "ERROR: %s failed with error code %s @ %s\n" \
                   "$2" "$3" "$(date +'%T')" | tee -a "$logfile"
            return 1
            ;;
        *)
            printf "CRASH: %s failed with error: %s @ %s. Script halted.\n" \
                   "$2" "$3" "$(date +'%T')" | tee -a "$logfile"
            exit 1
            ;;
    esac
}

display () {
    case "$1" in
        "title")
            printf "\n##### %s\n" "$2" | tee -a "$logfile"
            ;;
        "subtitle")
            printf "\n%s\n" "$2" | tee -a "$logfile"
            ;;
        "task")
            printf "...%s @ %s\n" "$2" "$(date +'%T')" | tee -a "$logfile"
            ;;
        *)
            printf "CRASH: display() called with incorrect arguments.\n"
            exit 1
            ;;
    esac
}

pacman_install() {
    action="installing $1 via pacman" && display "task" "$action"
    if ! (pacman --noconfirm --needed -S "$1" >> "$logfile" 2>&1); then
        action="retrying $1" && display "task" "$action"
        if ! (pacman --noconfirm --needed -S "$1" >> "$logfile" 2>&1); then
            action="retrying $1 once more" && display "task" "$action"
            (pacman --noconfirm --needed -S "$1" >> "$logfile" 2>&1) ||
                error "error" "$action" "$?"
        fi
    fi
}

aur_install() {
    action="installing $1 via the AUR" && display "task" "$action"
    if ! (sudo -u "$username" yay -S --noconfirm "$1" >> "$logfile" 2>&1); then
        action="retrying $1" && display "task" "$action"
        if ! (sudo -u "$username" yay -S --noconfirm "$1" >> "$logfile" 2>&1); then
            action="retrying $1 once more" && display "task" "$action"
            (sudo -u "$username" yay -S --noconfirm "$1" >> "$logfile" 2>&1) ||
                error "error" "$action" "$?"
        fi
    fi
}

git_install() {
    prog_name="$(basename "$1" .git)"
    build_dir="$source_dir/$prog_name"
    action="building & installing $prog_name from source"
    display "task" "$action"

    if ! (sudo -u "$username" git clone --depth 1 "$1" "$build_dir" >> "$logfile" 2>&1); then
        error "error" "cloning $prog_name - removing and retrying" "$?"
        rm -rf "$build_dir"
        (sudo -u "$username" git clone --depth 1 "$1" "$build_dir" >> "$logfile" 2>&1) || \
            error "crash" "re-cloning $prog_name" "$?"
    fi

    (cd "$build_dir" && make install >> "$logfile" 2>&1) || \
        error "error" "building $prog_name" "$?"
}

### ZFS-Specific Services
zfs_services() {
    display "title" "ZFS Services Verification"

    display "subtitle" "Verifying ZFS Snapshot Services"

    # Verify sanoid is configured and running
    action="checking sanoid configuration" && display "task" "$action"
    if [ -f /etc/sanoid/sanoid.conf ]; then
        echo "  Sanoid config found" | tee -a "$logfile"
    else
        warn "Sanoid config not found - snapshots may not be automatic"
    fi

    action="enabling sanoid timer" && display "task" "$action"
    systemctl enable --now sanoid.timer >> "$logfile" 2>&1 || \
        error "error" "$action" "$?"

    # Verify pacman hook
    action="checking pacman ZFS hook" && display "task" "$action"
    if [ -f /etc/pacman.d/hooks/zfs-snapshot.hook ]; then
        echo "  Pacman ZFS hook found" | tee -a "$logfile"
    else
        error "error" "Pacman ZFS hook not found" "1"
    fi

    # Show current ZFS status
    display "subtitle" "ZFS Pool Status"
    zpool status >> "$logfile" 2>&1
    zpool status

    display "subtitle" "ZFS Datasets"
    zfs list >> "$logfile" 2>&1
    zfs list

    display "subtitle" "Recent Snapshots"
    zfs list -t snapshot -o name,creation -s creation | tail -10
}

### Essential Services (ZFS variant - no btrfs tools)
essential_services() {
    display "title" "Essential Services"

    # ... [Keep all the same services from original archsetup EXCEPT:]
    # - Remove: timeshift-autosnap
    # - Remove: grub-btrfs
    # - Remove: grub-btrfsd configuration

    display "subtitle" "Randomness"
    pacman_install rng-tools
    systemctl enable rngd >> "$logfile" 2>&1

    display "subtitle" "Networking"
    pacman_install networkmanager

    display "subtitle" "Power"
    pacman_install upower
    systemctl enable upower >> "$logfile" 2>&1

    display "subtitle" "Secure Shell"
    pacman_install openssh
    systemctl enable sshd >> "$logfile" 2>&1

    display "subtitle" "Firewall"
    pacman_install ufw
    ufw default deny incoming >> "$logfile" 2>&1

    for protocol in \
        "80,443,8080/tcp" \
        "ssh" \
        "22000/tcp" "22000/udp" "21027/udp" \
        ; do
        action="adding ufw rule for $protocol" && display "task" "$action"
        ufw allow $protocol >> "$logfile" 2>&1 || error "error" "$action" "$?"
    done

    ufw limit 22/tcp >> "$logfile" 2>&1
    systemctl enable ufw.service >> "$logfile" 2>&1

    display "subtitle" "Network Service Discovery"
    pacman_install nss-mdns
    pacman_install avahi
    systemctl enable avahi-daemon.service >> "$logfile" 2>&1

    display "subtitle" "Job Scheduling"
    pacman_install cronie
    systemctl enable cronie >> "$logfile" 2>&1

    display "subtitle" "Package Repository Cache"
    pacman_install pacman-contrib
    systemctl enable --now paccache.timer >> "$logfile" 2>&1

    # ZFS snapshot services (already configured by install-archzfs)
    display "subtitle" "ZFS Snapshot Services"
    action="verifying sanoid timer" && display "task" "$action"
    systemctl enable --now sanoid.timer >> "$logfile" 2>&1 || \
        error "error" "$action" "$?"

    # NOTE: No grub-btrfs equivalent needed - ZFS boot environments
    # are handled differently (via GRUB entries or zfsbootmenu)
}

### Intro
intro() {
    printf "\n\nArchSetup-ZFS launched @ %s\n" "$(date +'%D %T')" | tee -a "$logfile"
    check_zfs

    STARTTIME=$(date +%s)
    errors_encountered=0

    [ -f "$logfile" ] && rm -f "$logfile"
    touch "$logfile"

    pacman -Q > "$packages_before" || error "crash" "generating package list" "$?"
}

### Outro
outro() {
    display "title" "Cleanup"

    action="forcing password change on first login" && display "task" "$action"
    chage -d 0 "$username" >> "$logfile" 2>&1 || error "error" "$action" "$?"

    display "subtitle" "Statistics"
    pacman -Q > "$packages_after"
    comm -13 --nocheck-order "$packages_before" "$packages_after" > "$archsetup_packages"

    ENDTIME=$(date +%s)
    totalsecs=$((ENDTIME - STARTTIME))
    mins=$((totalsecs / 60))
    secs=$((totalsecs % 60))
    new_packages=$(wc -l < "$archsetup_packages")

    printf "\n"
    printf "Completion time    : %s\n" "$(date +'%D %T')" | tee -a "$logfile"
    printf "Elapsed time       : %s minutes, %s seconds\n" "$mins" "$secs" | tee -a "$logfile"
    printf "Errors encountered : %s\n" "$errors_encountered" | tee -a "$logfile"
    printf "Log file location  : %s\n" "$logfile"
    printf "Packages installed : %s\n" "$new_packages"
    printf "\n"

    display "subtitle" "ZFS Quick Reference"
    printf "  List snapshots:     zfs list -t snapshot\n"
    printf "  Manual snapshot:    sudo zfs snapshot zroot/home@my-backup\n"
    printf "  Rollback:           sudo zfs rollback zroot/home@my-backup\n"
    printf "  Pool status:        zpool status\n"
    printf "  Replicate to NAS:   sudo zfs-replicate\n"
    printf "\n"
    printf "Please reboot before working with your new workstation.\n\n"

    printf "=== ARCHSETUP_ZFS_EXECUTION_COMPLETE ===\n" | tee -a "$logfile"
}

### Main - Placeholder for full implementation
main() {
    echo ""
    echo "=============================================="
    echo "  archsetup-zfs - ZFS Post-Installation Setup"
    echo "=============================================="
    echo ""
    echo "This is a skeleton script. The full implementation"
    echo "will include all sections from archsetup, modified"
    echo "for ZFS systems."
    echo ""
    echo "For now, run the original archsetup from:"
    echo "  ~/code/archsetup/archsetup"
    echo ""
    echo "The ZFS-specific services (sanoid, pacman hooks,"
    echo "syncoid) were already configured by install-archzfs."
    echo ""

    # When fully implemented, this will call:
    # run_step "intro" intro
    # run_step "prerequisites" prerequisites
    # run_step "user_customizations" user_customizations
    # run_step "aur_installer" aur_installer
    # run_step "essential_services" essential_services  # ZFS variant
    # run_step "xorg" xorg
    # run_step "dwm" dwm
    # run_step "desktop_environment" desktop_environment
    # run_step "developer_workstation" developer_workstation
    # run_step "supplemental_software" supplemental_software
    # run_step "boot_ux" boot_ux
    # run_step "zfs_services" zfs_services  # New ZFS-specific step
    # outro
}

main "$@"