#!/bin/bash
# Add an external disk for models. Run it through the one-liner (it asks for sudo once); no token, no reinstall,
# nothing is restarted, and it can be run again for the next disk. Any number of disks can be added.
#
#   curl -fsSL @@BASE@@/get | sh -s -- add-disk                   list the disks that can be added
#   curl -fsSL @@BASE@@/get | sh -s -- add-disk /dev/sdb1         add that partition: ext4, xfs or btrfs; its files are kept
#   curl -fsSL @@BASE@@/get | WXCH_CONFIRM=ERASE sh -s -- add-disk /dev/sdb1 --format
#                                                                  ERASE the partition, make it ext4 and add it
#   ... add-disk /dev/sdb1 --name backup                           choose the disk's id (default d1, d2, ...)
#
# The disk is mounted on demand at /mnt/wxch-ext/<id> (by its UUID, so the device name does not matter) and the models
# go into a folder named wxch-models on it; everything else on the disk is left alone. Only USB disks are accepted
# (WXCH_ALLOW_NON_USB=1 to allow others). The agent notices the new disk within seconds.
set -uo pipefail
BASE="${WXCH_BASE:?run this through the /get one-liner}"
WORK="$(mktemp -d /tmp/wxch-adddisk.XXXXXX)"
trap 'rm -rf "$WORK"' EXIT
curl -fsSL --connect-timeout 20 "$BASE/steps/lib.sh" -o "$WORK/lib.sh" || { echo "cannot download the helper script from $BASE" >&2; exit 1; }
bash -n "$WORK/lib.sh" || { echo "the helper script has a syntax error" >&2; exit 1; }
# shellcheck disable=SC1091
. "$WORK/lib.sh"
[ "$(id -u)" = 0 ] || die "must run as root: use the /get one-liner (it asks for sudo once)"
have blkid && have lsblk || die "lsblk and blkid are needed"

DEV=""; FORMAT=0; NAME=""
while [ $# -gt 0 ]; do
  case "$1" in
    --format) FORMAT=1 ;;
    --name) shift; NAME="${1:-}" ;;
    -h|--help) sed -n '2,/^set -uo/p' "$0" | sed '$d' | sed 's/^# \{0,1\}//'; exit 0 ;;
    -*) die "unknown option: $1 (see --help)" ;;
    *) [ -z "$DEV" ] || die "only one device at a time"; DEV="$1" ;;
  esac
  shift
done

registered() { python3 - "$1" <<'PY'
import json, os, sys
p = "/etc/wxch/node.json"
d = json.load(open(p)) if os.path.exists(p) else {}
for x in d.get("external_disks", []):
    if x.get("uuid") == sys.argv[1]:
        print(x.get("id", ""))
PY
}

# transport of the disk a device is on (usb, sata, nvme, ...)
tran_of() {
  pk="$(lsblk -no PKNAME "$1" 2>/dev/null | head -1 | tr -d ' ')"
  lsblk -dno TRAN "/dev/${pk:-${1#/dev/}}" 2>/dev/null | head -1 | tr -d ' '
}

# ---- no device: show what could be added ----
if [ -z "$DEV" ]; then
  echo "Disks that can be added (USB partitions):"
  found=0
  while read -r path typ; do
    fstype="$(lsblk -dno FSTYPE "$path" 2>/dev/null)"
    [ "$typ" = part ] || { [ "$typ" = disk ] && [ -n "$fstype" ]; } || continue
    [ "$(tran_of "$path")" = usb ] || [ "${WXCH_ALLOW_NON_USB:-}" = 1 ] || continue
    uuid="$(blkid -o value -s UUID "$path" 2>/dev/null)"
    note=""
    [ -n "$uuid" ] && [ -n "$(registered "$uuid")" ] && note="  (already added as $(registered "$uuid"))"
    case "$fstype" in
      ext4|xfs|btrfs) usable="its files are kept" ;;
      "") usable="no filesystem: needs --format" ;;
      *) usable="$fstype: needs --format (ERASES it)" ;;
    esac
    label="$(lsblk -dno LABEL "$path" 2>/dev/null)"
    printf '  %-12s %-6s %-8s %-16s %s%s\n' "$path" "${fstype:--}" "$(lsblk -dno SIZE "$path" | tr -d ' ')" "${label:--}" "$usable" "$note"
    mnt="$(lsblk -dno MOUNTPOINT "$path" 2>/dev/null | head -1)"
    [ -z "$mnt" ] || printf '               (mounted at %s)\n' "$mnt"
    found=1
  done < <(lsblk -rnp -o NAME,TYPE 2>/dev/null)
  [ "$found" = 1 ] || echo "  none found (plug the disk in; WXCH_ALLOW_NON_USB=1 also lists other disks)"
  echo
  echo "Add one with:  curl -fsSL $BASE/get | sh -s -- add-disk /dev/<partition>"
  exit 0
fi

# ---- checks ----
[ -b "$DEV" ] || die "$DEV is not a block device"
for m in $(lsblk -no MOUNTPOINT "$DEV" 2>/dev/null); do
  case "$m" in / | /boot | /boot/efi | /home | /usr | /var | /opt | "[SWAP]") die "refusing: $DEV holds $m (the system)" ;; esac
done
TRAN="$(tran_of "$DEV")"
[ "$TRAN" = usb ] || [ "${WXCH_ALLOW_NON_USB:-}" = 1 ] || die "refusing: $DEV is not on a USB disk (transport '${TRAN:-unknown}'); set WXCH_ALLOW_NON_USB=1 to allow it"

if [ -n "$NAME" ]; then
  echo "$NAME" | grep -Eq '^[a-z0-9][a-z0-9-]{0,9}$' || die "--name must be 1-10 characters: lowercase letters, digits, '-'"
fi

# ---- format (only when asked) ----
if [ "$FORMAT" = 1 ]; then
  [ "${WXCH_CONFIRM:-}" = ERASE ] || die "--format destroys everything on $DEV: also set WXCH_CONFIRM=ERASE"
  while read -r m; do
    [ -n "$m" ] || continue
    case "$m" in /media/* | /run/media/* | "$EXT_ROOT"/*) umount "$m" || die "cannot unmount $m" ;; *) die "$DEV is mounted at $m: unmount it first" ;; esac
  done < <(lsblk -no MOUNTPOINT "$DEV" 2>/dev/null)
  have mkfs.ext4 || die "mkfs.ext4 is not installed"
  say "formatting $DEV as ext4"
  mkfs.ext4 -F -L "WXCH-${NAME:-new}" "$DEV" >/dev/null || die "mkfs failed"
  partprobe "$DEV" >/dev/null 2>&1 || true
  sleep 1
fi

FSTYPE="$(blkid -o value -s TYPE "$DEV" 2>/dev/null)"
case "$FSTYPE" in
  ext4|xfs|btrfs) ;;
  "") die "$DEV has no filesystem. To use it, ERASE it and make it ext4: add --format and WXCH_CONFIRM=ERASE" ;;
  *) die "$DEV is $FSTYPE, which cannot keep Unix permissions, so it cannot hold models. Back it up, then add --format with WXCH_CONFIRM=ERASE (this erases it)" ;;
esac
UUID="$(blkid -o value -s UUID "$DEV" 2>/dev/null)"
[ -n "$UUID" ] || die "cannot read the UUID of $DEV"
FSLABEL="$(blkid -o value -s LABEL "$DEV" 2>/dev/null)"

# ---- already added? then just make sure it works ----
ID="$(registered "$UUID")"
if [ -n "$ID" ]; then
  say "this disk is already added as '$ID'; checking it"
else
  if [ -z "$NAME" ]; then
    n=1
    ids="$(python3 -c 'import json,os;p="/etc/wxch/node.json";d=json.load(open(p)) if os.path.exists(p) else {};print(" ".join(x.get("id","") for x in d.get("external_disks",[])))')"
    while case " $ids " in *" d$n "*) true ;; *) false ;; esac; do n=$((n+1)); done
    NAME="d$n"
  else
    case " $(python3 -c 'import json,os;p="/etc/wxch/node.json";d=json.load(open(p)) if os.path.exists(p) else {};print(" ".join(x.get("id","") for x in d.get("external_disks",[])))') " in
      *" $NAME "*) die "a disk called '$NAME' already exists: choose another --name" ;;
    esac
  fi
  ID="$NAME"
fi
MNT="$EXT_ROOT/$ID"

have systemctl && have systemd-escape || die "this needs systemd (systemctl and systemd-escape)"
mkdir -p "$MNT" "$EXT_ROOT"
if ! grep -qs "UUID=$UUID" /etc/fstab; then
  cp -p /etc/fstab "/etc/fstab.wxch-backup" 2>/dev/null || true
  echo "UUID=$UUID $MNT auto defaults,nofail,noauto,x-systemd.automount,x-systemd.idle-timeout=0,x-systemd.device-timeout=10 0 0" >> /etc/fstab
fi
install_ext_prep
UNIT="$(systemd-escape -p --suffix=mount "$MNT")"
INST="$(systemd-escape -p "$MNT")"
mkdir -p "/etc/systemd/system/$UNIT.wants"
ln -sf /etc/systemd/system/wxch-ext-prep@.service "/etc/systemd/system/$UNIT.wants/wxch-ext-prep@$INST.service"
systemctl daemon-reload

mountpoint -q "$MNT" || mount "$MNT" || die "could not mount $DEV at $MNT (see: journalctl -xe)"
/usr/local/sbin/wxch-ext-prep "$MNT" "$OWNER"
[ -e "$MNT/wxch-models/.wxch-models" ] || die "the disk was mounted but could not be prepared (see: journalctl -u 'wxch-ext-prep@*')"

# ---- a disk from the first version of this feature: one disk, found by its label WXCH-MODELS, mounted at
# ---- /mnt/wxch-external, with the marker and the runtimes' folders at the root of the disk. Bring it over: renames on the
# ---- same disk (instant), the machine's links re-pointed, the old mount and its fstab line retired.
OLD="${WXCH_OLD_EXT_DIR:-/mnt/wxch-external}"
if [ -e "$MNT/.wxch-models" ]; then
  say "this disk has the older layout: moving its model folders into wxch-models/"
  S="$MNT/wxch-models"
  for d in ollama vllm; do
    [ -d "$MNT/$d" ] || continue
    rmdir "$S/$d/blobs" "$S/$d" 2>/dev/null || true      # the empty folders that were just created
    if [ -e "$S/$d" ]; then warn "$S/$d already has files: leaving $MNT/$d where it is"; continue; fi
    mv "$MNT/$d" "$S/$d" || warn "could not move $MNT/$d"
  done
  /usr/local/sbin/wxch-ext-prep "$MNT" "$OWNER"          # recreates any folder that is missing
  rm -f "$MNT/.wxch-models"
  relink() { # folder maxdepth: links that pointed into the old mount now point into the new place
    [ -d "$1" ] || return 0
    find "$1" -maxdepth "$2" -type l 2>/dev/null | while read -r l; do
      t="$(readlink "$l")"
      case "$t" in
        "$OLD"/*) n="$S/${t#"$OLD"/}"; if [ -e "$n" ]; then ln -s "$n" "$l.wxch-new" && mv -T "$l.wxch-new" "$l"; fi ;;
      esac
    done
  }
  relink "$MODELS_DIR/ollama/blobs" 1
  relink "$MODELS_DIR/vllm" 2
  if grep -qs "^LABEL=[^ ]*[[:space:]]\{1,\}${OLD}[[:space:]]" /etc/fstab; then
    sed -i "s|^\(LABEL=[^ ]*[[:space:]]\{1,\}${OLD}[[:space:]].*\)|# replaced by add-disk (kept for reference): \1|" /etc/fstab
  fi
  systemctl stop "$(systemd-escape -p --suffix=automount "$OLD")" "$(systemd-escape -p --suffix=mount "$OLD")" >/dev/null 2>&1 || true
  mountpoint -q "$OLD" && { umount "$OLD" 2>/dev/null || warn "could not unmount the old mount $OLD (harmless; it goes away at the next reboot)"; }
  rm -f /etc/systemd/system/wxch-ext-prep.service /etc/systemd/system/*.mount.wants/wxch-ext-prep.service
  systemctl daemon-reload
  say "the older layout is migrated; the machine's links to its models were re-pointed"
fi

# the agent reads this file again every few seconds: no restart needed
python3 - "$ID" "$UUID" "${FSLABEL:-$ID}" "$EXT_ROOT" <<'PY'
import json, os, sys
ident, uuid, label, root = sys.argv[1:5]
p = "/etc/wxch/node.json"
os.makedirs("/etc/wxch", exist_ok=True)
d = json.load(open(p)) if os.path.exists(p) else {}
d["external_root"] = root
disks = [x for x in d.get("external_disks", []) if x.get("uuid") != uuid]
disks.append({"id": ident, "uuid": uuid, "label": label})
d["external_disks"] = sorted(disks, key=lambda x: x["id"])
json.dump(d, open(p, "w"), indent=2)
open(p, "a").write("\n")
PY

grant_ollama_blobs
allow_runner_root "$EXT_ROOT"

say "disk '$ID' is added: $DEV ($FSTYPE) mounted at $MNT, models go to $MNT/wxch-models"
df -h "$MNT" | tail -1
say "nothing was restarted. The agent notices the disk within a few seconds; the Models page shows it. Add more disks the same way."