# /etc/profile.d/nimux-first-login.sh
# Nimux first-login bootstrap (mirror) — numbered disk selection + safeguards + default proposal

case $- in *i*) ;; *) return 0 ;; esac

SENTINEL="/var/lib/nimux/firstboot.done"
[ -f "$SENTINEL" ] && return 0

run_bootstrap() {
  set -euo pipefail

  need() { command -v "$1" >/dev/null 2>&1 || { echo "Missing: $1"; exit 1; }; }
  need zpool; need zfs; need lsblk; command -v blkid >/dev/null 2>&1 || true; command -v blockdev >/dev/null 2>&1 || true

  CUR_HOST="$(hostname)"
  POOL="${CUR_HOST}-zfs"
  FSTAB="/etc/fstab"
  OVERLAY_ROOT="/mnt/overlays/root"
  ADMINUSER=""
  NEW_LANG="${LANG:-en_US.UTF-8}"
  DISK1=""; DISK2=""
  DEFAULT_IDX1=""; DEFAULT_IDX2=""

  # -------------------- Helpers --------------------

  canonical() { readlink -f "$1"; }

  has_partitions() {
    # returns 0 (true) if the disk has child partitions
    local d="$1"
    lsblk -n "$d" -o TYPE | grep -q '^part$'
  }

  fstype_of() {
    # Try lsblk first, then blkid
    local d="$1" fs
    fs="$(lsblk -dn -o FSTYPE "$d" 2>/dev/null | tr -d ' ')"
    if [ -z "$fs" ] && command -v blkid >/dev/null 2>&1; then
      fs="$(blkid -o value -s TYPE "$d" 2>/dev/null | tr -d ' ')"
    fi
    echo "${fs:-}"
  }

  bytes_of() {
    # size in bytes (fallback to lsblk if blockdev not present)
    local d="$1"
    if command -v blockdev >/dev/null 2>&1; then
      blockdev --getsize64 "$d" 2>/dev/null || echo 0
    else
      lsblk -bdn -o SIZE "$d" 2>/dev/null || echo 0
    fi
  }

  is_whole_disk() {
    # TYPE=disk
    local d="$1"
    [ "$(lsblk -dn -o TYPE "$d" 2>/dev/null)" = "disk" ]
  }

  is_zfs_member() {
    local d="$1"
    [ "$(fstype_of "$d")" = "zfs_member" ]
  }

  print_disks_table() {
    echo "Available whole disks:"
    printf "  %-5s %-20s %-12s %-12s %s\n" "No." "PATH" "SIZE" "FSTYPE" "MODEL"
    local i=1
    while read -r name size type fstype model; do
      [ "$type" = "disk" ] || continue
      local path="/dev/$name"
      local sz="$size"
      printf "  [%d]   %-20s %-12s %-12s %s\n" "$i" "$path" "$sz" "${fstype:--}" "${model:-"-"}"
      i=$((i+1))
    done < <(lsblk -dn -o NAME,SIZE,TYPE,FSTYPE,MODEL)
  }

  build_disk_arrays() {
    # Build arrays for selection UI
    mapfile -t DISK_LINES < <(lsblk -dn -o NAME,SIZE,TYPE,FSTYPE,MODEL | awk '$3=="disk"{print}')
    mapfile -t DISK_PATHS < <(printf "%s\n" "${DISK_LINES[@]}" | awk '{print "/dev/"$1}')
    [ "${#DISK_PATHS[@]}" -ge 1 ] || { echo "No whole disks found."; exit 1; }
  }

  propose_two_largest_empty() {
    # Choose two largest disks that are: TYPE=disk, no partitions, no FSTYPE, not zfs_member
    local candidates=()
    local sizes=()

    for p in "${DISK_PATHS[@]}"; do
      is_whole_disk "$p" || continue
      if has_partitions "$p"; then continue; fi
      local fs; fs="$(fstype_of "$p")"
      [ -n "$fs" ] && continue
      is_zfs_member "$p" && continue
      local b; b="$(bytes_of "$p")"
      candidates+=("$p"); sizes+=("$b")
    done

    if [ "${#candidates[@]}" -lt 2 ]; then
      # Not enough clean disks to propose
      DEFAULT_IDX1=""; DEFAULT_IDX2=""
      return 0
    fi

    # Sort candidates by size desc (simple bubble for small lists)
    local idxs=()
    local i; for ((i=0;i<${#candidates[@]};i++)); do idxs+=("$i"); done
    local swapped=1 j tmp
    while [ $swapped -eq 1 ]; do
      swapped=0
      for ((j=0;j<${#idxs[@]}-1;j++)); do
        if [ "${sizes[${idxs[$j]}]}" -lt "${sizes[${idxs[$j+1]}]}" ]; then
          tmp="${idxs[$j]}"; idxs[$j]="${idxs[$j+1]}"; idxs[$j+1]="$tmp"; swapped=1
        fi
      done
    done

    local top1="${candidates[${idxs[0]}]}"
    local top2="${candidates[${idxs[1]}]}"

    # Map to numbered list indices shown to user
    local n i=1
    for n in "${DISK_PATHS[@]}"; do
      if [ "$(canonical "$n")" = "$(canonical "$top1")" ]; then DEFAULT_IDX1="$i"; fi
      if [ "$(canonical "$n")" = "$(canonical "$top2")" ]; then DEFAULT_IDX2="$i"; fi
      i=$((i+1))
    done
  }

  warn_if_risky() {
    # Return 0 if OK; non-zero if user declines risk
    local d="$1" label="$2"
    local risk=0 msg=""
    if has_partitions "$d"; then risk=1; msg+="\n - $label has existing partitions"; fi
    local fs; fs="$(fstype_of "$d")"
    if [ -n "$fs" ]; then risk=1; msg+="\n - $label has existing filesystem: $fs"; fi
    if [ "$risk" -eq 0 ]; then return 0; fi
    echo -e "WARNING:$msg"
    read -r -p "Type 'force' to override and continue with $label: " force
    [ "${force:-}" = "force" ] || { echo "Refusing risky selection for $label."; return 1; }
    return 0
  }

  pick_two_disks() {
    echo
    print_disks_table
    build_disk_arrays
    propose_two_largest_empty

    local default_hint=""
    if [ -n "$DEFAULT_IDX1" ] && [ -n "$DEFAULT_IDX2" ]; then
      default_hint=" [${DEFAULT_IDX1} ${DEFAULT_IDX2}]"
      echo
      echo "Proposed (largest clean disks): ${DEFAULT_IDX1} and ${DEFAULT_IDX2}"
      echo "Press Enter to accept, or type two numbers to override."
    fi

    while :; do
      echo
      read -r -p "Select TWO disks by number (e.g. '1 3')${default_hint}: " CHOICE
      if [ -z "${CHOICE// }" ] && [ -n "$DEFAULT_IDX1" ] && [ -n "$DEFAULT_IDX2" ]; then
        idx1="$DEFAULT_IDX1"; idx2="$DEFAULT_IDX2"
      else
        CHOICE="$(echo "$CHOICE" | tr ',' ' ' | xargs)"
        set +e; set -- $CHOICE; set -e
        if [ $# -ne 2 ]; then echo "Please provide exactly two numbers."; continue; fi
        idx1="$1"; idx2="$2"
      fi

      case "$idx1" in (*[!0-9]*|'') echo "Invalid first selection."; continue;; esac
      case "$idx2" in (*[!0-9]*|'') echo "Invalid second selection."; continue;; esac
      [ "$idx1" -ge 1 ] && [ "$idx1" -le "${#DISK_PATHS[@]}" ] || { echo "First number out of range."; continue; }
      [ "$idx2" -ge 1 ] && [ "$idx2" -le "${#DISK_PATHS[@]}" ] || { echo "Second number out of range."; continue; }
      [ "$idx1" -ne "$idx2" ] || { echo "Disks must be different."; continue; }

      local cand1="${DISK_PATHS[$((idx1-1))]}"
      local cand2="${DISK_PATHS[$((idx2-1))]}"

      # Risk checks
      warn_if_risky "$cand1" "Disk 1" || continue
      warn_if_risky "$cand2" "Disk 2" || continue

      DISK1="$cand1"; DISK2="$cand2"
      echo
      echo "Selected:"
      echo "  Disk 1: $DISK1"
      echo "  Disk 2: $DISK2"
      read -r -p "Proceed to create MIRRORED pool on these disks? [y/N]: " yn
      case "${yn:-N}" in y|Y) break;; *) echo "Re-selecting…";; esac
    done
  }

  create_pool_and_datasets() {
    echo
    echo "== Creating mirrored ZFS pool and datasets =="
    if zpool list -H -o name 2>/dev/null | grep -qx "$POOL"; then
      echo "Pool '$POOL' already exists; skipping creation."
    else
      ASHIFT="12"
      zpool create -f \
        -o ashift="$ASHIFT" \
        -O compression=zstd -O atime=off -O xattr=sa -O acltype=posixacl \
        "$POOL" mirror "$DISK1" "$DISK2"
    fi

    zfs set mountpoint=none "$POOL" || true

    create_ds() {  # $1=dataset $2=mountpoint|"none"
      local ds="$1" mp="$2"
      if ! zfs list -H -o name | grep -qx "$ds"; then
        if [ "$mp" = "none" ]; then
          zfs create -o mountpoint=none "$ds"
        else
          zfs create -o mountpoint="$mp" "$ds"
        fi
      else
        [ "$mp" = "none" ] && zfs set mountpoint=none "$ds" || zfs set mountpoint="$mp" "$ds"
      fi
    }

    create_ds "${POOL}"                         "none"
    create_ds "${POOL}/data"                    "/mnt/data"
    create_ds "${POOL}/docker"                  "/var/lib/docker"
    create_ds "${POOL}/overlay-layers"          "none"
    create_ds "${POOL}/overlay-layers/root"     "$OVERLAY_ROOT"

    mkdir -p /mnt/data /var/lib/docker "$OVERLAY_ROOT"
    zfs mount -a
  }

  setup_overlays() {
    echo
    echo "== Setting up overlays for /etc and /home =="
    mkdir -p "${OVERLAY_ROOT}/etc" "${OVERLAY_ROOT}/etc.work" \
             "${OVERLAY_ROOT}/home" "${OVERLAY_ROOT}/home.work"
    chmod 700 "${OVERLAY_ROOT}/etc.work" "${OVERLAY_ROOT}/home.work"
    touch "$FSTAB"

    add_fstab_line() { local line="$1" pattern="$2"; grep -qsE "$pattern" "$FSTAB" || echo "$line" >> "$FSTAB"; }

    add_fstab_line \
      "ZFS=${POOL}/overlay-layers/root  $OVERLAY_ROOT  zfs  defaults,nofail  0  0" \
      "^ZFS=${POOL}/overlay-layers/root[[:space:]]+$OVERLAY_ROOT[[:space:]]+zfs"

    add_fstab_line \
      "overlay  /etc   overlay  nofail,x-systemd.requires-mounts-for=$OVERLAY_ROOT,lowerdir=/etc,upperdir=${OVERLAY_ROOT}/etc,workdir=${OVERLAY_ROOT}/etc.work  0  0" \
      "^overlay[[:space:]]+/etc[[:space:]]+overlay.*upperdir=${OVERLAY_ROOT}/etc,workdir=${OVERLAY_ROOT}/etc.work"

    add_fstab_line \
      "overlay  /home  overlay  nofail,x-systemd.requires-mounts-for=$OVERLAY_ROOT,lowerdir=/home,upperdir=${OVERLAY_ROOT}/home,workdir=${OVERLAY_ROOT}/home.work  0  0" \
      "^overlay[[:space:]]+/home[[:space:]]+overlay.*upperdir=${OVERLAY_ROOT}/home,workdir=${OVERLAY_ROOT}/home.work"
  }

  ask_hostname_and_maybe_rename_pool() {
    echo
    echo "== Hostname =="
    read -r -p "New hostname (leave blank to keep '$CUR_HOST'): " NEW_HOSTNAME
    NEW_HOSTNAME="${NEW_HOSTNAME:-$CUR_HOST}"

    if command -v hostnamectl >/dev/null 2>&1; then
      hostnamectl set-hostname "$NEW_HOSTNAME"
    else
      echo "$NEW_HOSTNAME" > /etc/hostname
    fi

    local NEWPOOL="${NEW_HOSTNAME}-zfs"
    if [ "$NEWPOOL" != "$POOL" ]; then
      echo "Renaming pool '$POOL' -> '$NEWPOOL' via export/import…"
      zfs unmount -a || true
      sed -i "s/ZFS=${POOL}\//ZFS=${NEWPOOL}\//g" "$FSTAB"
      zpool export "$POOL"
      zpool import "$POOL" "$NEWPOOL"
      zfs set mountpoint=none "$NEWPOOL" || true
      zfs set mountpoint=/mnt/data            "$NEWPOOL/data" || true
      zfs set mountpoint=/var/lib/docker      "$NEWPOOL/docker" || true
      zfs set mountpoint=none                 "$NEWPOOL/overlay-layers" || true
      zfs set mountpoint="$OVERLAY_ROOT"      "$NEWPOOL/overlay-layers/root" || true
      zfs mount -a
      POOL="$NEWPOOL"
    fi
  }

  set_locale_step() {
    echo
    echo "== Locale =="
    read -r -p "Locale (LANG) [${NEW_LANG}]: " LANG_INPUT
    NEW_LANG="${LANG_INPUT:-$NEW_LANG}"
    printf 'LANG=%s\n' "$NEW_LANG" > /etc/locale.conf
    if command -v localectl >/dev/null 2>&1; then
      localectl set-locale "LANG=$NEW_LANG" || true
    fi
  }

  create_admin_user() {
    echo
    echo "== Create admin user with sudo rights =="
    while :; do
      read -r -p "Admin username (e.g., nimuxadmin): " ADMINUSER
      [ -n "${ADMINUSER:-}" ] && break
      echo "Username cannot be empty."
    done

    if id "$ADMINUSER" >/dev/null 2>&1; then
      echo "User '$ADMINUSER' already exists."
    else
      if [ -x /bin/bash ]; then SHELLPATH=/bin/bash
      elif [ -x /bin/sh ]; then SHELLPATH=/bin/sh
      else SHELLPATH=/bin/sh; fi
      useradd -m -s "$SHELLPATH" "$ADMINUSER"
      echo "Set password for $ADMINUSER:"
      passwd "$ADMINUSER"
    fi

    if command -v sudo >/dev/null 2>&1; then
      getent group sudo >/dev/null 2>&1 || groupadd sudo
      usermod -aG sudo "$ADMINUSER" || true
      mkdir -p /etc/sudoers.d
      SUDO_FILE="/etc/sudoers.d/90-${ADMINUSER}"
      [ -f "$SUDO_FILE" ] || { echo "${ADMINUSER} ALL=(ALL) ALL" > "$SUDO_FILE"; chmod 440 "$SUDO_FILE"; }
    else
      echo "NOTE: 'sudo' not installed; grant temporary root via 'su' until sudo is available."
    fi
  }

  disable_root_login() {
    echo
    echo "== Disable root logins =="
    passwd -l root || true
    if [ -f /etc/ssh/sshd_config ]; then
      if grep -qE '^\s*PermitRootLogin\s+' /etc/ssh/sshd_config; then
        sed -i 's/^\s*PermitRootLogin\s\+.*/PermitRootLogin no/' /etc/ssh/sshd_config
      else
        printf '\nPermitRootLogin no\n' >> /etc/ssh/sshd_config
      fi
      if grep -qE '^\s*PermitEmptyPasswords\s+' /etc/ssh/sshd_config'; then
        sed -i 's/^\s*PermitEmptyPasswords\s\+.*/PermitEmptyPasswords no/' /etc/ssh/sshd_config
      else
        printf 'PermitEmptyPasswords no\n' >> /etc/ssh/sshd_config
      fi
      (command -v systemctl >/dev/null 2>&1 && systemctl reload sshd) || \
      (command -v rc-service >/dev/null 2>&1 && rc-service sshd reload) || true
    fi
  }

  mark_completion() {
    mkdir -p "$(dirname "$SENTINEL")"
    echo "ok" > "$SENTINEL"
    echo
    echo "Bootstrap complete."
    echo " - Pool: ${POOL}"
    echo " - Overlays for /etc and /home will take effect on next reboot."
    echo " - Locale: ${NEW_LANG}"
    echo " - Admin user: ${ADMINUSER}"
    echo " - Root login disabled"
  }

  # -------------------- Ordered flow --------------------
  echo
  echo "=== Nimux First-Login ==="
  build_disk_arrays
  pick_two_disks
  create_pool_and_datasets
  setup_overlays
  ask_hostname_and_maybe_rename_pool
  set_locale_step
  create_admin_user
  disable_root_login
  mark_completion
}

if [ "$EUID" -ne 0 ]; then
  echo
  echo "[Nimux] First-boot setup needs elevated privileges. Re-running with sudo…"
  echo
  sudo bash -c "$(declare -f run_bootstrap); run_bootstrap" && \
    echo "Setup completed (or already done)." || \
    echo "Setup aborted or failed."
else
  run_bootstrap
fi

[ -f "$SENTINEL" ] && chmod -x /etc/profile.d/nimux-first-login.sh 2>/dev/null || true