#!/usr/bin/env bash
# prepare_mariadb_usb.sh
#
# From myrddin: identify a USB block device, wipe/partition/format it for MariaDB,
# mount it temporarily, and drop a 1-click setup.sh on the root of the USB filesystem
# that (when run on the target host) creates an "ansible" user and installs an SSH public key.
#
# SAFETY: This script can DESTROY DATA. Read prompts carefully.

set -euo pipefail

USB_LABEL="MARIADB"
MOUNT_BASE="/mnt/usb-setup"
TARGET_MOUNTPOINT="/srv/mariadb"
ANSIBLE_USER="ansible"
SUDO_NOPASSWD="yes"

usage() {
  cat <<'EOF'
Usage:
  sudo ./prepare_mariadb_usb.sh --pubkey-file /path/to/ansible.pub
  sudo ./prepare_mariadb_usb.sh --pubkey "ssh-ed25519 AAAA... comment"

Options:
  --pubkey-file PATH   Path to SSH public key to install for the 'ansible' user (recommended)
  --pubkey STRING      SSH public key string (quoted)
  --device /dev/sdX    Skip auto-detect and use this device (DANGEROUS)
  --label LABEL        Filesystem label (default: MARIADB)
  --mount DIR          Temp mountpoint base (default: /mnt/usb-setup)
  --target-mount DIR   Target host mountpoint (setup.sh will mount here; default: /srv/mariadb)
  --help               Show this help

What it does:
  1) Picks a removable USB disk
  2) Wipes signatures, creates GPT + single ext4 partition
  3) Labels it and mounts it
  4) Writes /setup.sh on the USB root to create 'ansible' user + install authorized_keys
  5) The generated setup.sh can also add an /etc/fstab entry and mount the USB on the target host
EOF
}

PUBKEY=""
PUBKEY_FILE=""
DEVICE=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    --pubkey-file) PUBKEY_FILE="${2:-}"; shift 2 ;;
    --pubkey) PUBKEY="${2:-}"; shift 2 ;;
    --device) DEVICE="${2:-}"; shift 2 ;;
    --label) USB_LABEL="${2:-}"; shift 2 ;;
    --mount) MOUNT_BASE="${2:-}"; shift 2 ;;
    --target-mount) TARGET_MOUNTPOINT="${2:-}"; shift 2 ;;
    --help|-h) usage; exit 0 ;;
    *) echo "Unknown arg: $1"; usage; exit 1 ;;
  esac
done

if [[ -n "$PUBKEY_FILE" ]]; then
  if [[ ! -f "$PUBKEY_FILE" ]]; then
    echo "ERROR: --pubkey-file '$PUBKEY_FILE' not found"
    exit 1
  fi
  PUBKEY="$(cat "$PUBKEY_FILE")"
fi

if [[ -z "$PUBKEY" ]]; then
  echo "ERROR: You must provide --pubkey-file or --pubkey"
  exit 1
fi

if [[ "$EUID" -ne 0 ]]; then
  echo "ERROR: run as root (sudo)"
  exit 1
fi

need_cmd() { command -v "$1" >/dev/null 2>&1 || { echo "ERROR: missing '$1'"; exit 1; }; }
need_cmd lsblk
need_cmd wipefs
need_cmd sgdisk
need_cmd mkfs.ext4
need_cmd mount
need_cmd umount
need_cmd udevadm
need_cmd partprobe

# ---- Device selection ----
# Auto-detect: usb + removable + disk
if [[ -z "$DEVICE" ]]; then
  echo "Scanning for candidate removable USB disks…"
  mapfile -t CANDS < <(
    lsblk -dn -o NAME,TYPE,TRAN,RM,SIZE,MODEL | awk '
      $2=="disk" && $3=="usb" && $4=="1" {
        print $1 "|" $5 "|" $6
      }'
  )

  if [[ ${#CANDS[@]} -eq 0 ]]; then
    echo "ERROR: No ~4GB removable USB disks found."
    echo "Hint: plug the USB stick in, then run: lsblk -o NAME,SIZE,TRAN,RM,MODEL"
    exit 1
  fi

  echo "Candidates:"
  i=1
  for c in "${CANDS[@]}"; do
    name="${c%%|*}"
    rest="${c#*|}"
    echo "  [$i] /dev/${name}  (${rest})"
    ((i++))
  done

  if [[ ${#CANDS[@]} -eq 1 ]]; then
    DEVICE="/dev/${CANDS[0]%%|*}"
  else
    read -r -p "Select device number to FORMAT: " sel
    [[ "$sel" =~ ^[0-9]+$ ]] || { echo "Invalid selection"; exit 1; }
    (( sel>=1 && sel<=${#CANDS[@]} )) || { echo "Out of range"; exit 1; }
    DEVICE="/dev/${CANDS[$((sel-1))]%%|*}"
  fi
fi

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

echo
echo "Selected device: $DEVICE"
lsblk -o NAME,SIZE,TYPE,TRAN,RM,MODEL "$DEVICE" || true
echo

# Refuse to operate on root disk (best-effort check)
ROOT_SRC="$(findmnt -n -o SOURCE /)"
ROOT_DEV=""
if [[ "$ROOT_SRC" == /dev/* ]]; then
  ROOT_PARENT="$(lsblk -no PKNAME "$ROOT_SRC" 2>/dev/null | head -n 1 || true)"
  if [[ -n "$ROOT_PARENT" ]]; then
    ROOT_DEV="/dev/${ROOT_PARENT}"
  else
    # Fallback (handles /dev/sda2 -> /dev/sda)
    ROOT_DEV="$(echo "$ROOT_SRC" | sed -E 's/[0-9]+$//')"
  fi
fi

if [[ -n "$ROOT_DEV" && "$DEVICE" == "$ROOT_DEV" ]]; then
  echo "ERROR: Refusing to operate on root device '$DEVICE'"
  exit 1
fi

read -r -p "TYPE 'FORMAT' to wipe and format $DEVICE: " confirm
if [[ "$confirm" != "FORMAT" ]]; then
  echo "Aborted."
  exit 1
fi

# ---- Wipe / partition / format ----
echo "Unmounting any mounted partitions on $DEVICE…"
# shellcheck disable=SC2012
for p in $(lsblk -ln -o NAME "$DEVICE" | tail -n +2); do
  mp="$(findmnt -n -o TARGET "/dev/$p" 2>/dev/null || true)"
  if [[ -n "$mp" ]]; then
    echo "  umount /dev/$p ($mp)"
    umount "/dev/$p"
  fi
done

echo "Wiping filesystem signatures…"
wipefs -a "$DEVICE"

echo "Creating GPT + single partition…"
sgdisk --zap-all "$DEVICE"
sgdisk -n 1:0:0 -t 1:8300 -c 1:"${USB_LABEL}" "$DEVICE"

echo "Informing kernel of partition changes…"
udevadm settle
partprobe "$DEVICE" 2>/dev/null || true
udevadm settle

PART="${DEVICE}1"
# NVMe and some devices use p1 naming
if [[ ! -b "$PART" ]]; then
  PART="${DEVICE}p1"
fi
if [[ ! -b "$PART" ]]; then
  echo "ERROR: Partition device not found after partitioning ($DEVICE)."
  lsblk "$DEVICE"
  exit 1
fi

echo "Formatting $PART as ext4 with label '$USB_LABEL'…"
mkfs.ext4 -F -L "$USB_LABEL" "$PART"

# ---- Mount and write bootstrap ----
mkdir -p "$MOUNT_BASE"
MNT="$(mktemp -d -p "$MOUNT_BASE" "${USB_LABEL}.XXXX")"
echo "Mounting $PART at $MNT…"
mount "$PART" "$MNT"

# Data dir hint for MariaDB
mkdir -p "$MNT/mariadb"
chmod 700 "$MNT/mariadb"

# Write setup.sh onto the USB root
cat > "$MNT/setup.sh" <<EOF
#!/usr/bin/env bash
# setup.sh (runs on TARGET HOST, not on myrddin)
# Creates '${ANSIBLE_USER}' user, installs SSH key, grants sudo if configured.
set -euo pipefail

ANSIBLE_USER="${ANSIBLE_USER}"
SUDO_NOPASSWD="${SUDO_NOPASSWD}"
USB_LABEL="${USB_LABEL}"
MARIADB_MOUNTPOINT="${TARGET_MOUNTPOINT}"
MARIADB_FSTAB_OPTS="noatime,nofail,x-systemd.device-timeout=10"

PUBKEY='${PUBKEY}'

if [[ "\$EUID" -ne 0 ]]; then
  echo "Run as root: sudo ./setup.sh"
  exit 1
fi

if ! id "\$ANSIBLE_USER" >/dev/null 2>&1; then
  echo "Creating user \$ANSIBLE_USER…"
  useradd -m -s /bin/bash "\$ANSIBLE_USER"
else
  echo "User \$ANSIBLE_USER already exists."
fi

echo "Installing SSH authorized_keys for \$ANSIBLE_USER…"
install -d -m 700 -o "\$ANSIBLE_USER" -g "\$ANSIBLE_USER" "/home/\$ANSIBLE_USER/.ssh"
printf "%s\n" "\$PUBKEY" > "/home/\$ANSIBLE_USER/.ssh/authorized_keys"
chown "\$ANSIBLE_USER:\$ANSIBLE_USER" "/home/\$ANSIBLE_USER/.ssh/authorized_keys"
chmod 600 "/home/\$ANSIBLE_USER/.ssh/authorized_keys"

if [[ "\$SUDO_NOPASSWD" == "yes" ]]; then
  echo "Granting passwordless sudo to \$ANSIBLE_USER…"
  cat > "/etc/sudoers.d/90-\$ANSIBLE_USER" <<SUD
\$ANSIBLE_USER ALL=(ALL) NOPASSWD:ALL
SUD
  chmod 440 "/etc/sudoers.d/90-\$ANSIBLE_USER"
fi

echo "Configuring MariaDB USB mount (label=\$USB_LABEL, mountpoint=\$MARIADB_MOUNTPOINT)…"

if ! command -v blkid >/dev/null 2>&1; then
  echo "ERROR: missing 'blkid' on target host. Install util-linux and re-run."
  exit 1
fi

if ! blkid -L "\$USB_LABEL" >/dev/null 2>&1; then
  echo "ERROR: USB filesystem with label '\$USB_LABEL' not found."
  echo "Hint: lsblk -f"
  exit 1
fi

install -d -m 0755 "\$MARIADB_MOUNTPOINT"

FSTAB_LINE="LABEL=\$USB_LABEL \$MARIADB_MOUNTPOINT ext4 \$MARIADB_FSTAB_OPTS 0 2"
if ! grep -qs "^LABEL=\$USB_LABEL[[:space:]]\+\$MARIADB_MOUNTPOINT[[:space:]]" /etc/fstab; then
  echo "Adding fstab entry for MariaDB USB…"
  printf "%s\n" "\$FSTAB_LINE" >> /etc/fstab
else
  echo "fstab entry already present."
fi

echo "Mounting \$MARIADB_MOUNTPOINT…"
mount "\$MARIADB_MOUNTPOINT" || mount -a

mkdir -p "\$MARIADB_MOUNTPOINT/mariadb"
chmod 700 "\$MARIADB_MOUNTPOINT/mariadb"

echo
echo "Done. Test from control machine:"
echo "  ssh \$ANSIBLE_USER@<host>"
EOF

chmod +x "$MNT/setup.sh"

cat > "$MNT/README.txt" <<EOF
This USB stick was prepared for:
- MariaDB data directory (on target host): ${TARGET_MOUNTPOINT}/mariadb
- Ansible bootstrap: /setup.sh

On the target host:
  1) Plug in USB
  2) Mount it temporarily somewhere (example): sudo mount /dev/sdX1 /mnt
  3) Run: sudo /mnt/setup.sh
     (setup.sh will add an /etc/fstab entry using LABEL=${USB_LABEL} and mount it at ${TARGET_MOUNTPOINT})
Then you can run Ansible against that host using the installed SSH key.

Filesystem label: ${USB_LABEL}
EOF

sync
echo
echo "USB prepared successfully."
echo "Mounted at: $MNT"
echo "Contents:"
ls -la "$MNT"
echo
echo "Next steps:"
echo "  umount '$MNT' && rmdir '$MNT'"
echo "  (then move the USB stick to the target host and run setup.sh there)"