#!/usr/bin/env bash
# make-usb.sh
# 
# Write a NixOS installer ISO to a USB stick (raw dd)
# and optionally append a writable "extras" partition with files
# (this repo, SSH keys, etc.) that you want available during install.

set -euo pipefail



###   Styling   ###

if [[ -t 1 ]]; then
    GREEN=$'\e[32m'; YELLOW=$'\e[33m'; RED=$'\e[31m'; BLUE=$'\e[34m'
    BOLD=$'\e[1m'; RESET=$'\e[0m'
else
    GREEN=""; YELLOW=""; RED=""; BLUE=""; BOLD=""; RESET=""
fi
log()   { printf '%s[+]%s %s\n'  "$GREEN"  "$RESET" "$*"; }
warn()  { printf '%s[!]%s %s\n'  "$YELLOW" "$RESET" "$*" >&2; }
err()   { printf '%s[-]%s %s\n'  "$RED"    "$RESET" "$*" >&2; }
info()  { printf '%s[i]%s %s\n'  "$BLUE"   "$RESET" "$*"; }
die()   { err "$1"; exit "${2:-1}"; }

ISO_PATH=""
DEVICE=""
SKIP_EXTRAS=0
ALLOW_INTERNAL=0
ASSUME_YES=0
REPO_ROOT=""

while (($#)); do
    case "$1" in
        --iso)            ISO_PATH="${2:?missing path}"; shift 2 ;;
        --device|--dev)   DEVICE="${2:?missing path}"; shift 2 ;;
        --no-extras)      SKIP_EXTRAS=1; shift ;;
        --allow-internal) ALLOW_INTERNAL=1; shift ;;
        --yes|-y)         ASSUME_YES=1; shift ;;
        --repo-root)      REPO_ROOT="${2:?missing path}"; shift 2 ;;
        -h|--help)
            sed -n '2,16p' "$0" | sed 's/^# \{0,1\}//'
            exit 0 ;;
        *) die "unknown arg: $1" 2 ;;
    esac
done



###   Preflight checks   ###

(( EUID == 0 )) || die "run as root: sudo $0 $*" 1

need_tty=0
[[ -z "$DEVICE" || -z "$ISO_PATH" ]] && need_tty=1
(( ASSUME_YES )) || need_tty=1
if (( need_tty )) && ! [[ -t 0 ]]; then
    die "stdin is not a tty; pass --iso, --device, and --yes for non-interactive use" 2
fi

# Required commands. We try to be helpful about missing ones on NixOS.
REQUIRED=(lsblk dd sync wipefs sgdisk partprobe blockdev mount umount mountpoint mkfs.exfat awk)
missing=()
for cmd in "${REQUIRED[@]}"; do
    command -v "$cmd" >/dev/null 2>&1 || missing+=("$cmd")
done
if ((${#missing[@]})); then
    err "missing tools: ${missing[*]}"
    cat >&2 <<'EOF'

On NixOS / with Nix installed, re-run inside a shell that has them:
  nix shell nixpkgs#util-linux nixpkgs#coreutils nixpkgs#gptfdisk \
            nixpkgs#exfatprogs nixpkgs#parted -c sudo ./scripts/make-usb.sh

On Debian/Ubuntu:
  sudo apt install util-linux coreutils gdisk exfatprogs parted
EOF
    exit 3
fi

# Locate repo root if not provided (used as a copy candidate for extras).
if [[ -z "$REPO_ROOT" ]]; then
    if REPO_ROOT="$(git -C "$(dirname "$0")" rev-parse --show-toplevel 2>/dev/null)"; then
        :
    else
        REPO_ROOT=""
    fi
fi



###   Cleanup trap   ###

TMP_MNT=""
cleanup() {
    set +e
    if [[ -n "$TMP_MNT" && -d "$TMP_MNT" ]]; then
        mountpoint -q "$TMP_MNT" && umount "$TMP_MNT"
        rmdir "$TMP_MNT" 2>/dev/null
    fi
}
trap cleanup EXIT INT TERM



###   Helpers   ###

human() { numfmt --to=iec --suffix=B "$1" 2>/dev/null || echo "$1"; }

is_removable() {
    # /sys/block/<name>/removable: 1 = removable (USB stick, SD via USB reader),
    # 0 = fixed (internal disk, many "USB" enclosures lie though).
    local name; name="$(basename "$1")"
    [[ "$(cat "/sys/block/$name/removable" 2>/dev/null || echo 0)" == "1" ]]
}

contains_root_or_boot() {
    # Refuse if any partition on this disk holds /, /boot, or active swap.
    local dev="$1"
    lsblk -lnpo MOUNTPOINT,NAME "$dev" | awk -v d="$dev" '
        $1=="/" || $1=="/boot" || $1=="[SWAP]" { found=1 }
        END { exit !found }
    '
}

select_device() {
    log "Available whole disks:"
    printf '   %-3s %-14s %-8s %-6s %-16s %s\n' "Idx" "Device" "Size" "Bus" "Removable" "Model"

    # Use lsblk's pairs mode for robust parsing (handles spaces in MODEL).
    # lsblk -P emits lines of KEY="value" KEY="value" ... we sanity-check
    # the line matches that strict shape before eval'ing it.
    local -a CANDS=()
    local line NAME SIZE MODEL TRAN RM TYPE
    while IFS= read -r line; do
        [[ "$line" =~ ^([A-Z]+=\"[^\"]*\"[[:space:]]*)+$ ]] || continue
        NAME=""; SIZE=""; MODEL=""; TRAN=""; RM=""; TYPE=""
        eval "$line"
        [[ "${TYPE:-}" == "disk" ]] || continue
        # Skip optical, loop, ram, zram, dm.
        case "$NAME" in
            /dev/sr*|/dev/loop*|/dev/ram*|/dev/zram*|/dev/dm-*) continue ;;
        esac
        CANDS+=("$NAME"$'\t'"${SIZE:-?}"$'\t'"${TRAN:-?}"$'\t'"${RM:-0}"$'\t'"${MODEL:-?}")
    done < <(lsblk -dpP -o NAME,SIZE,MODEL,TRAN,RM,TYPE)

    ((${#CANDS[@]})) || die "no candidate disks found" 4

    local i=1
    for c in "${CANDS[@]}"; do
        IFS=$'\t' read -r n s t r m <<<"$c"
        local marker=""
        if [[ "$r" == "1" ]]; then marker="${GREEN}yes${RESET}"; else marker="${YELLOW}no${RESET}"; fi
        if contains_root_or_boot "$n"; then
            marker="${RED}SYSTEM DISK${RESET}"
        fi
        printf '   %-3s %-14s %-8s %-6s %-16b %s\n' "$i" "$n" "$s" "$t" "$marker" "$m"
        ((i++))
    done
    echo
    echo "Enter 0 to abort."
    echo

    local idx
    while :; do
        read -r -p "Select device [1-${#CANDS[@]}]: " idx
        [[ "$idx" =~ ^[0-9]+$ ]] || { warn "not a number"; continue; }
        (( idx == 0 )) && die "aborted by user" 1
        (( idx >= 1 && idx <= ${#CANDS[@]} )) || { warn "out of range"; continue; }
        IFS=$'\t' read -r n _ _ r _ <<<"${CANDS[idx-1]}"
        if contains_root_or_boot "$n"; then
            err "$n is the running system disk — refusing"; continue
        fi
        if [[ "$r" != "1" ]] && (( ! ALLOW_INTERNAL )); then
            err "$n is not removable; pass --allow-internal if you really mean it"
            continue
        fi
        DEVICE="$n"
        return 0
    done
}



###   Pick device   ###

if [[ -z "$DEVICE" ]]; then
    select_device
else
    [[ -b "$DEVICE" ]] || die "$DEVICE is not a block device" 2
    # Reject partitions. We want whole disks only.
    if [[ "$(lsblk -dno TYPE "$DEVICE" 2>/dev/null)" != "disk" ]]; then
        die "$DEVICE is not a whole disk (partition or other type)" 4
    fi
    if contains_root_or_boot "$DEVICE"; then
        die "$DEVICE holds /, /boot, or active swap — refusing" 4
    fi
    if ! is_removable "$DEVICE" && (( ! ALLOW_INTERNAL )); then
        die "$DEVICE is not removable; pass --allow-internal to override" 4
    fi
fi

DEV_BYTES="$(blockdev --getsize64 "$DEVICE")"
log "Target: $DEVICE ($(human "$DEV_BYTES"))"



###   Pick ISO   ###

if [[ -z "$ISO_PATH" ]]; then
    while :; do
        read -r -e -p "Path to NixOS ISO: " ISO_PATH
        ISO_PATH="${ISO_PATH/#\~/$HOME}"
        [[ -f "$ISO_PATH" ]] && break
        warn "not a file: $ISO_PATH"
    done
fi
[[ -f "$ISO_PATH" ]] || die "ISO not found: $ISO_PATH" 2

ISO_BYTES="$(stat -c '%s' "$ISO_PATH")"
log "ISO: $ISO_PATH ($(human "$ISO_BYTES"))"

# Sanity: file should look like an ISO 9660 image.
if command -v file >/dev/null 2>&1; then
    if ! file -b "$ISO_PATH" | grep -qiE 'iso 9660|udf'; then
        warn "$(file -b "$ISO_PATH")"
        warn "file(1) doesn't recognise this as ISO 9660. Continuing anyway."
    fi
fi

# Sanity: ISO must fit on disk with room for an extras partition.
MIN_EXTRAS_MB=64
if (( ISO_BYTES + MIN_EXTRAS_MB*1024*1024 > DEV_BYTES )); then
    if (( ISO_BYTES > DEV_BYTES )); then
        die "ISO is larger than the disk" 4
    fi
    warn "Less than ${MIN_EXTRAS_MB}MiB free after ISO; extras partition will be skipped."
    SKIP_EXTRAS=1
fi

# Optional sha256 sidecar verification.
if [[ -f "$ISO_PATH.sha256" ]]; then
    log "Verifying $ISO_PATH.sha256..."
    ( cd "$(dirname "$ISO_PATH")" && sha256sum -c "$(basename "$ISO_PATH").sha256" ) \
        || die "sha256 verification failed" 5
fi



###   Final confirmation   ###

cat <<EOF

${BOLD}About to:${RESET}
  ${RED}WIPE${RESET}            $DEVICE  ($(human "$DEV_BYTES"))
  Write ISO       $ISO_PATH
  Add extras      $( ((SKIP_EXTRAS)) && echo "no" || echo "yes (exFAT, in trailing free space)" )

${YELLOW}All data on $DEVICE will be destroyed.${RESET}
EOF

if (( ! ASSUME_YES )); then
    read -r -p "Type the device path to confirm (e.g. $DEVICE): " typed
    [[ "$typed" == "$DEVICE" ]] || die "confirmation didn't match — aborting" 1
fi



###   Unmount anything on the target   ###

log "Unmounting any partitions on $DEVICE..."
while read -r part mnt; do
    [[ -n "$mnt" ]] || continue
    warn "umount $part ($mnt)"
    umount "$part" 2>/dev/null || umount -l "$part" 2>/dev/null || warn "could not umount $part"
done < <(lsblk -lnpo NAME,MOUNTPOINT "$DEVICE" | awk '$2!=""')

###   Wipe & write   ###
log "Wiping signatures..."
wipefs -a "$DEVICE" >/dev/null

log "Writing ISO with dd (this can take a few minutes)..."
dd if="$ISO_PATH" of="$DEVICE" bs=4M status=progress conv=fsync oflag=direct \
    || die "dd failed" 5
sync
log "ISO written."

partprobe "$DEVICE" 2>/dev/null || true
sleep 1



###   Extras partition   ###

if (( SKIP_EXTRAS )); then
    log "Skipping extras partition."
else
    log "Preparing extras partition in trailing free space..."

    # The NixOS ISO uses an isohybrid GPT. The backup GPT header still sits at
    # the *original ISO end*, not the actual disk end. `sgdisk -e` relocates it
    # so we can add a partition in the freed trailing region.
    # `sgdisk -n` with 0:0:0 creates a new partition starting at the first free
    # sector and ending at the last free sector (i.e., filling the rest).
    sgdisk -e "$DEVICE" >/dev/null
    sgdisk -n 0:0:0 -t 0:0700 -c 0:"EXTRAS" "$DEVICE" >/dev/null
    partprobe "$DEVICE" 2>/dev/null || true
    sleep 1

    # Find the new partition (highest-numbered one on the disk).
    EXTRAS_PART="$(lsblk -lnpo NAME,TYPE "$DEVICE" \
        | awk '$2=="part"{print $1}' | tail -n1)"
    [[ -b "$EXTRAS_PART" ]] || die "extras partition didn't appear" 5
    log "Extras partition: $EXTRAS_PART"

    log "Formatting $EXTRAS_PART as exFAT..."
    mkfs.exfat -L EXTRAS "$EXTRAS_PART" >/dev/null 2>&1 \
        || die "mkfs.exfat failed" 5

    TMP_MNT="$(mktemp -d)"
    mount "$EXTRAS_PART" "$TMP_MNT"

    echo
    log "What to copy onto the extras partition? (space-separated numbers; Enter for none)"
    echo
    echo "   1) This repo                ${REPO_ROOT:-<not detected>}"
    echo "   2) Your public SSH keys     ~/.ssh/*.pub"
    echo "   3) An arbitrary file/dir    (you'll be prompted for the path)"
    echo "   4) A custom message/README  (you type it inline)"
    echo

    read -r -p "Selection: " sel || sel=""
    for tok in $sel; do
        case "$tok" in
            1)
                if [[ -z "$REPO_ROOT" ]]; then
                    warn "no repo detected; pass --repo-root or run from a git checkout"
                    continue
                fi
                log "Copying repo -> /repo (git tracked files only)..."
                # Use `git archive` if possible, it respects .gitignore and skips .git.
                if git -C "$REPO_ROOT" rev-parse >/dev/null 2>&1; then
                    mkdir -p "$TMP_MNT/repo"
                    git -C "$REPO_ROOT" archive --format=tar HEAD \
                        | tar -x -C "$TMP_MNT/repo"
                else
                    mkdir -p "$TMP_MNT/repo"
                    cp -aT "$REPO_ROOT" "$TMP_MNT/repo"
                fi
                ;;
            2)
                # Run as the invoking user, not root, so ~ resolves correctly.
                user_home="$(getent passwd "${SUDO_USER:-$USER}" | cut -d: -f6)"
                if compgen -G "$user_home/.ssh/*.pub" >/dev/null; then
                    mkdir -p "$TMP_MNT/ssh-keys"
                    cp "$user_home"/.ssh/*.pub "$TMP_MNT/ssh-keys/"
                    log "Copied $(ls "$user_home"/.ssh/*.pub | wc -l) public key(s)."
                else
                    warn "no .pub keys found in $user_home/.ssh"
                fi
                ;;
            3)
                read -r -e -p "Path to copy: " extra_path
                extra_path="${extra_path/#\~/$HOME}"
                if [[ -e "$extra_path" ]]; then
                    cp -a "$extra_path" "$TMP_MNT/"
                    log "Copied $extra_path"
                else
                    warn "not found: $extra_path"
                fi
                ;;
            4)
                info "Type your note. End with a line containing only EOF:"
                : >"$TMP_MNT/NOTE.txt"
                while IFS= read -r line; do
                    [[ "$line" == "EOF" ]] && break
                    printf '%s\n' "$line" >>"$TMP_MNT/NOTE.txt"
                done
                ;;
            *) warn "unknown option: $tok" ;;
        esac
    done

    sync
    umount "$TMP_MNT"
    rmdir "$TMP_MNT"
    TMP_MNT=""
fi



###   Done   ###

log "Final sync..."
sync
blockdev --flushbufs "$DEVICE" 2>/dev/null || true

echo
log "${BOLD}Done.${RESET} You can unplug $DEVICE now."
echo
info "Boot the new machine from this USB. Once at the installer shell:"
info "  - If you copied the repo, it's on the EXTRAS partition (mountable as exFAT)."
info "  - To install: partition the target disk, mount at /mnt, then"
info "      nixos-install --flake /mnt/extras/repo#<hostname>"
echo