#!/bin/bash
# Runs INSIDE a privileged container (see run-add-disk.sh): the add-disk script against real (loop) disks.
# systemd is not running here, so systemctl is a stub; mounting, fstab, UUIDs, file ownership and the config file are real.
set -uo pipefail
PASS=0; FAIL=0
ok()   { echo "  PASS  $1"; PASS=$((PASS+1)); }
bad()  { echo "  FAIL  $1"; FAIL=$((FAIL+1)); }
check() { if eval "$2"; then ok "$1"; else bad "$1"; tail -4 "$(ls -t /tmp/o* 2>/dev/null | head -1)" 2>/dev/null | sed "s/^/        | /"; fi; }

export DEBIAN_FRONTEND=noninteractive
apt-get update -qq >/dev/null 2>&1
apt-get install -y -qq e2fsprogs util-linux python3 curl systemd acl xfsprogs dosfstools >/dev/null 2>&1 || { echo "cannot install the tools"; exit 1; }
useradd -m tester 2>/dev/null
printf '#!/bin/sh\nexit 0\n' > /usr/local/bin/systemctl; chmod +x /usr/local/bin/systemctl
mkdir -p /etc/wxch /mnt/models/ollama/blobs
echo '{"models_dir": "/mnt/models/ollama"}' > /etc/wxch/node.json
echo '{"image": "x", "models_dir": "/mnt/models/vllm", "agent_uid": 1000}' > /etc/wxch/runner.json

# loop device nodes are not created on demand inside a container
for i in $(seq 0 63); do [ -e /dev/loop$i ] || mknod /dev/loop$i b 7 $i 2>/dev/null; done
cleanup() { # detach the loop devices of this test (mk runs in a subshell, so look them up by their image files)
  for m in /mnt/wxch-ext/*; do umount "$m" 2>/dev/null; done
  for img in /tmp/disk1.img /tmp/disk2.img /tmp/fat.img /tmp/plain.img /tmp/oldlayout.img; do
    for d in $(losetup -j "$img" -n -O NAME 2>/dev/null); do umount "$d" 2>/dev/null; losetup -d "$d" 2>/dev/null; done
  done
}
trap cleanup EXIT

mk() { # name fs -> prints the loop device
  truncate -s 300M "/tmp/$1.img"
  dev="$(losetup -f --show "/tmp/$1.img")" || { echo "no loop device" >&2; exit 1; }
  case "$2" in
    ext4) mkfs.ext4 -q -F -L "$1" "$dev" >/dev/null 2>&1 ;;
    xfs) mkfs.xfs -q -f "$dev" >/dev/null 2>&1 ;;
    vfat) mkfs.vfat "$dev" >/dev/null 2>&1 ;;
  esac
  echo "$dev"
}
D1="$(mk disk1 ext4)"; D2="$(mk disk2 xfs)"; DV="$(mk fat vfat)"; DN="$(mk plain none)"
echo "devices: $D1 $D2 $DV $DN"

ADD() { WXCH_BASE="file:///t/files" WXCH_USER=tester WXCH_ALLOW_NON_USB=1 bash /t/files/add-disk.sh "$@"; }
cfg() { python3 -c "import json;d=json.load(open('/etc/wxch/node.json'));$1"; }

echo "== 1. add the first disk (ext4, with a file on it that must survive)"
mkdir -p /mnt/x && mount "$D1" /mnt/x && echo precious > /mnt/x/my-photos.txt && umount /mnt/x
ADD "$D1" > /tmp/o1 2>&1; rc=$?
check "exits 0" "[ $rc = 0 ]"
check "it says the disk is added" "grep -q \"disk 'd1' is added\" /tmp/o1"
check "mounted at /mnt/wxch-ext/d1" "mountpoint -q /mnt/wxch-ext/d1"
check "marker and folders exist, owned by the agent user" "[ -e /mnt/wxch-ext/d1/wxch-models/.wxch-models ] && [ \"\$(stat -c %U /mnt/wxch-ext/d1/wxch-models/ollama/blobs)\" = tester ] && [ -d /mnt/wxch-ext/d1/wxch-models/vllm ]"
check "the disk's own files are untouched" "[ \"\$(cat /mnt/wxch-ext/d1/my-photos.txt)\" = precious ]"
check "registered in node.json" "[ \"\$(cfg 'print(d[\"external_disks\"][0][\"id\"], d[\"external_root\"])')\" = 'd1 /mnt/wxch-ext' ]"
check "fstab entry by UUID with automount" "grep -q 'UUID=.* /mnt/wxch-ext/d1 auto .*x-systemd.automount' /etc/fstab"
check "the runner may use the external root, other settings kept" "[ \"\$(python3 -c \"import json;c=json.load(open('/etc/wxch/runner.json'));print(c['extra_roots'], c['agent_uid'])\")\" = \"['/mnt/wxch-ext'] 1000\" ]"
check "the prep helper and its unit template are installed" "[ -x /usr/local/sbin/wxch-ext-prep ] && grep -q 'ConditionPathIsMountPoint=%f' /etc/systemd/system/wxch-ext-prep@.service"
check "the ACL lets the agent user write in Ollama's blobs" "getfacl -p /mnt/models/ollama/blobs 2>/dev/null | grep -q 'user:tester:rwx'"

echo "== 2. the same disk again is recognised, not added twice"
ADD "$D1" > /tmp/o2 2>&1; rc=$?
check "exits 0" "[ $rc = 0 ]"
check "says it is already added" "grep -q 'already added as' /tmp/o2"
check "still one disk, one fstab line" "[ \"\$(cfg 'print(len(d[\"external_disks\"]))')\" = 1 ] && [ \"\$(grep -c 'wxch-ext/d1' /etc/fstab)\" = 1 ]"

echo "== 3. a second disk (xfs) gets the next id"
ADD "$D2" > /tmp/o3 2>&1; rc=$?
check "exits 0" "[ $rc = 0 ]"
check "it is d2 and mounted" "mountpoint -q /mnt/wxch-ext/d2 && [ -e /mnt/wxch-ext/d2/wxch-models/.wxch-models ]"
check "both disks are in node.json, runner.json not duplicated" "[ \"\$(cfg 'print(\",\".join(x[\"id\"] for x in d[\"external_disks\"]))')\" = 'd1,d2' ] && [ \"\$(grep -c wxch-ext /etc/wxch/runner.json)\" = 1 ]"

echo "== 4. refusals"
ADD "$DV" > /tmp/o4 2>&1; rc=$?
check "a vfat disk is refused without touching it" "[ $rc != 0 ] && grep -q 'cannot keep Unix permissions' /tmp/o4 && ! grep -q vfat /etc/fstab"
ADD "$DN" > /tmp/o5 2>&1; rc=$?
check "a disk without a filesystem is refused" "[ $rc != 0 ] && grep -q 'has no filesystem' /tmp/o5"
ADD "$DV" --format > /tmp/o6 2>&1; rc=$?
check "--format needs WXCH_CONFIRM=ERASE" "[ $rc != 0 ] && grep -q 'WXCH_CONFIRM=ERASE' /tmp/o6"
ADD /dev/null > /tmp/o7 2>&1; rc=$?
check "not a block device" "[ $rc != 0 ]"
WXCH_BASE="file:///t/files" WXCH_USER=tester bash /t/files/add-disk.sh "$D1" > /tmp/o8 2>&1; rc=$?
check "a non-USB disk is refused unless allowed" "[ $rc != 0 ] || grep -q 'already added' /tmp/o8"
ADD "$DN" --name 'Bad Name' > /tmp/o9 2>&1; rc=$?
check "a bad --name is refused" "[ $rc != 0 ] && grep -q -- '--name must be' /tmp/o9"

echo "== 5. --format makes an empty disk usable (this is the only way anything is erased)"
WXCH_CONFIRM=ERASE ADD "$DN" --format --name scratch > /tmp/o10 2>&1; rc=$?
check "exits 0 and the disk is added as 'scratch'" "[ $rc = 0 ] && mountpoint -q /mnt/wxch-ext/scratch && [ -e /mnt/wxch-ext/scratch/wxch-models/.wxch-models ]"
WXCH_CONFIRM=ERASE ADD "$DV" --format > /tmp/o11 2>&1; rc=$?
check "a vfat disk can be erased on request and becomes d3" "[ $rc = 0 ] && mountpoint -q /mnt/wxch-ext/d3"

echo "== 7. a disk from the first version (label WXCH-MODELS, folders at its root, mounted at /mnt/wxch-external)"
DO="$(mk oldlayout ext4)"; e2label "$DO" WXCH-MODELS
mkdir -p /mnt/wxch-external && mount "$DO" /mnt/wxch-external
mkdir -p /mnt/wxch-external/ollama/blobs /mnt/wxch-external/vllm/org/repo /mnt/models/vllm/org
: > /mnt/wxch-external/.wxch-models
echo blobdata > /mnt/wxch-external/ollama/blobs/sha256-aaa
echo '{"models":{"m:1b":["sha256:aaa"]}}' > /mnt/wxch-external/ollama/index.json
echo '{}' > /mnt/wxch-external/vllm/org/repo/config.json
ln -s /mnt/wxch-external/ollama/blobs/sha256-aaa /mnt/models/ollama/blobs/sha256-aaa
ln -s /mnt/wxch-external/vllm/org/repo /mnt/models/vllm/org/repo
echo 'LABEL=WXCH-MODELS /mnt/wxch-external auto defaults,nofail,noauto,x-systemd.automount 0 0' >> /etc/fstab
ADD "$DO" > /tmp/o13 2>&1; rc=$?
UUIDO="$(blkid -o value -s UUID "$DO")"
NEWID="$(cfg "print([x['id'] for x in d['external_disks'] if x['uuid']=='$UUIDO'][0])")"
check "the old disk is added" "[ $rc = 0 ] && grep -q 'older layout' /tmp/o13 && mountpoint -q /mnt/wxch-ext/$NEWID"
check "its folders moved into wxch-models and the old marker is gone" "[ -e /mnt/wxch-ext/$NEWID/wxch-models/ollama/blobs/sha256-aaa ] && [ -e /mnt/wxch-ext/$NEWID/wxch-models/ollama/index.json ] && [ -e /mnt/wxch-ext/$NEWID/wxch-models/vllm/org/repo/config.json ] && [ ! -e /mnt/wxch-ext/$NEWID/.wxch-models ] && [ ! -e /mnt/wxch-ext/$NEWID/ollama ]"
check "the machine's Ollama link points at the new place and still reads" "[ \"\$(cat /mnt/models/ollama/blobs/sha256-aaa)\" = blobdata ] && readlink /mnt/models/ollama/blobs/sha256-aaa | grep -q '^/mnt/wxch-ext/$NEWID/wxch-models/ollama/blobs/'"
check "the machine's vLLM link points at the new place and still reads" "[ -e /mnt/models/vllm/org/repo/config.json ] && readlink /mnt/models/vllm/org/repo | grep -q '^/mnt/wxch-ext/$NEWID/wxch-models/vllm/'"
check "the old fstab line is retired and the old mount is gone" "! grep -q '^LABEL=WXCH-MODELS' /etc/fstab && grep -q '# replaced by add-disk' /etc/fstab && ! mountpoint -q /mnt/wxch-external"

echo "== 6. the system is never offered"
ROOTDEV="$(findmnt -n -o SOURCE / | head -1)"
WXCH_ALLOW_NON_USB=1 WXCH_BASE="file:///t/files" WXCH_USER=tester bash /t/files/add-disk.sh "$ROOTDEV" > /tmp/o12 2>&1; rc=$?
check "the disk holding / is refused" "[ $rc != 0 ]"

echo
echo "passed $PASS, failed $FAIL"
[ "$FAIL" = 0 ] && echo "ALL ADD-DISK CHECKS PASSED"
exit "$FAIL"