#!/usr/bin/env bash
set -euo pipefail

DISK="${1:-}"
FLAKE_DIR="$(dirname "$(realpath "$0")")"

if [[ -z "$DISK" ]]; then
    echo "Usage: $0 <disk_device>"
    echo "Example: $0 /dev/nvme0n1"
    echo "         $0 /dev/sda"
    exit 1
fi

if [[ ! -b "$DISK" ]]; then
    echo "Error: $DISK is not a block device"
    exit 1
fi

echo "Installing NixOS to $DISK using flake at $FLAKE_DIR"

# Create temporary directories
mkdir -p /tmp/secureboot
mkdir -p /tmp/persistent

# Function to cleanup on exit
cleanup() {
    echo "Cleaning up..."
    rm -rf /tmp/secureboot 2>/dev/null || true
    rm -rf /tmp/persistent 2>/dev/null || true
}
trap cleanup EXIT

# Decrypt secureboot keys using the key in the repo
echo "Decrypting secureboot keys..."
if [[ ! -f "$FLAKE_DIR/usb-secrets/usb-secrets-key" ]]; then
    echo "Error: usb-secrets-key not found at $FLAKE_DIR/usb-secrets/usb-secrets-key"
    exit 1
fi

nix-shell -p age --run "age -d -i '$FLAKE_DIR/usb-secrets/usb-secrets-key' '$FLAKE_DIR/secrets/secureboot.tar.age'" | \
    tar -x -C /tmp/secureboot

echo "Secureboot keys extracted"

# Extract persistent partition secrets
echo "Extracting persistent partition contents..."
if [[ -f "$FLAKE_DIR/secrets/persistent.tar" ]]; then
    tar -xzf "$FLAKE_DIR/secrets/persistent.tar" -C /tmp/persistent
    echo "Persistent partition contents extracted"
else
    echo "Warning: persistent.tar not found, skipping persistent secrets"
fi

# Check if disko-install is available
if ! command -v disko-install >/dev/null 2>&1; then
    echo "Running disko-install via nix..."
    DISKO_INSTALL="nix run github:nix-community/disko#disko-install --"
else
    DISKO_INSTALL="disko-install"
fi

echo "Running disko-install to partition, format, and install NixOS..."

# Build the extra-files arguments
EXTRA_FILES_ARGS=(
    --extra-files /tmp/secureboot /etc/secureboot
    --extra-files "$FLAKE_DIR/usb-secrets/usb-secrets-key" /mnt/usb-secrets/usb-secrets-key
)

# Add each top-level item from persistent separately to avoid nesting
# cp -ar creates /dst/src when copying directories, so we need to copy each item
#
# Also disko-install actually copies the files from extra-files, so we are good here
if [[ -d /tmp/persistent ]] && [[ -n "$(ls -A /tmp/persistent 2>/dev/null)" ]]; then
    for item in /tmp/persistent/*; do
        if [[ -e "$item" ]]; then
            basename=$(basename "$item")
            EXTRA_FILES_ARGS+=(--extra-files "$item" "/persistent/$basename")
        fi
    done
fi

# Run disko-install with secureboot keys available
sudo $DISKO_INSTALL \
    --mode format \
    --flake "$FLAKE_DIR#muffin" \
    --disk main "$DISK" \
    "${EXTRA_FILES_ARGS[@]}"