#!/usr/bin/env bash
# Write a NixOS installer ISO to a USB stick (raw dd).
#
# The script self-elevates via sudo if not already running as root, so
# you can run it as your normal user.
#
# After booting the resulting USB on the target machine, fetch this repo
# over the network with `git clone <url>` from the installer shell.

set -euo pipefail



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=""
ALLOW_INTERNAL=0
ASSUME_YES=0

# Save original args before the parser shifts them away; needed if we
# later re-exec under sudo.
ORIG_ARGS=("$@")

while (($#)); do
    case "$1" in
        --iso)            ISO_PATH="${2:?missing path}"; shift 2 ;;
        --device|--dev)   DEVICE="${2:?missing path}"; shift 2 ;;
        --allow-internal) ALLOW_INTERNAL=1; shift ;;
        --yes|-y)         ASSUME_YES=1; shift ;;
        -h|--help)
            awk '/^#!/ {next} /^[^#]/ {exit} {sub(/^# ?/, ""); print}' "$0"
            exit 0 ;;
        *) die "unknown arg: $1" 2 ;;
    esac
done



# Self-elevate via sudo if not already root. We re-exec rather than tell
# the user to prefix the command, so `just make-usb` and `./scripts/make-usb.sh`
# both Just Work.
if (( EUID != 0 )); then
    if ! command -v sudo >/dev/null 2>&1; then
        die "must run as root and sudo is not installed" 1
    fi
    # ORIG_ARGS may be empty; ${ORIG_ARGS[@]+"${ORIG_ARGS[@]}"} expands
    # safely under set -u even when the array is unset.
    exec sudo --preserve-env=PATH -- "$0" ${ORIG_ARGS[@]+"${ORIG_ARGS[@]}"}
fi

# Require an interactive tty for any prompt we still need.
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.
REQUIRED=(lsblk dd sync wipefs blockdev mount umount 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 -c sudo ./scripts/make-usb.sh

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



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
}



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"))"



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.
if (( ISO_BYTES > DEV_BYTES )); then
    die "ISO is larger than the disk" 4
fi

# Optional sha256 sidecar verification.
# Resolve through symlinks so `make-usb --iso iso/latest-...iso` finds
# the sidecar that lives next to the real file.
ISO_REAL="$(readlink -f "$ISO_PATH")"
if [[ -f "$ISO_REAL.sha256" ]]; then
    log "Verifying $ISO_REAL.sha256..."
    ( cd "$(dirname "$ISO_REAL")" && sha256sum -c "$(basename "$ISO_REAL").sha256" ) \
        || die "sha256 verification failed" 5
elif [[ -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



cat <<EOF

${BOLD}About to:${RESET}
  ${RED}WIPE${RESET}            $DEVICE  ($(human "$DEV_BYTES"))
  Write ISO       $ISO_PATH

${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



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!=""')



log "Wiping signatures..."
wipefs -a "$DEVICE" >/dev/null

log "Writing ISO with dd (this can take a few minutes)..."
if ! dd if="$ISO_PATH" of="$DEVICE" bs=4M status=progress conv=fsync oflag=direct; then
    err "dd failed."
    # Surface the most recent kernel complaints about this device so the
    # user can distinguish bad-block / hardware failure ("Medium Error",
    # "Hardware Error") from cable/port flakiness ("usb ... reset",
    # "Communication failure") without having to dig through dmesg.
    dev_short="$(basename "$DEVICE")"
    if dmesg_out="$(dmesg 2>/dev/null | grep -E "(${dev_short}|usb [0-9]+-[0-9]+)" | tail -10)" \
       && [[ -n "$dmesg_out" ]]; then
        warn "Recent kernel messages mentioning ${dev_short} or USB events:"
        printf '%s\n' "$dmesg_out" | sed 's/^/    /' >&2
        warn "Look for 'Medium Error' / 'Hardware Error' (bad flash → replace stick),"
        warn "or 'reset' / 'Communication failure' (cable/port → try a different one)."
    else
        warn "Run \`sudo dmesg | tail -30\` to see why the kernel rejected the write."
    fi
    exit 5
fi

log "Final sync (flushing kernel buffers; may take a moment)..."
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 "  - Connect to the network (Ethernet, or 'sudo systemctl start wpa_supplicant' + 'wpa_cli', or 'nmtui' on the graphical ISO)."
info "  - Clone your repo:   git clone <your-repo-url> /tmp/notfiles"
info "  - Partition the target disk, mount at /mnt, then:"
info "      nixos-install --flake /tmp/notfiles#<hostname>"
echo