#!/usr/bin/env bash
# Prepare a flexible N-disk RAID0 builder stripe in Hetzner rescue mode.
# This is destructive: it wipes every disk selected in the saved layout plan.

set -euo pipefail

REPO_ROOT="$(git rev-parse --show-toplevel)"
source "$REPO_ROOT/tools/lib/builder-common.sh"
cd "$TOTALLY_SPIES_REPO_ROOT"

PLAN_FILE="${PLAN_FILE:-$TOTALLY_SPIES_BUILDER_ROOTFS_PLAN_FILE}"
STATE_FILE="${STATE_FILE:-$TOTALLY_SPIES_BUILDER_HOST_STATE_FILE}"
SSH_KEY="${SSH_KEY:-$HOME/.ssh/id_ed25519}"
SERVER_HOST="${1:-}"

if [ -z "$SERVER_HOST" ]; then
  SERVER_HOST="$(builder_current_host)"
fi

if [ -z "$SERVER_HOST" ]; then
  echo "Usage: prepare_rescue_stripe.sh <server-host>" >&2
  echo "Or save server_host in $STATE_FILE first via the rescue probe task." >&2
  exit 1
fi

builder_require_file "$PLAN_FILE"
builder_require_file "$SSH_KEY"

run_remote() {
  builder_ssh root "$SERVER_HOST" "$@"
}

remote_payload="$(python3 - <<'PY' "$PLAN_FILE"
import json, sys
from pathlib import Path
plan = json.loads(Path(sys.argv[1]).read_text())
out = {
    'md_device': plan['md_device'],
    'mount_point': plan['mount_point'],
    'fs_type': plan['fs_type'],
    'members': [
        {
            'disk': m['disk'],
        }
        for m in plan['members']
    ],
}
print(json.dumps(out))
PY
)"
payload_b64="$(printf '%s' "$remote_payload" | base64 -w0)"

echo "Preparing rescue stripe on $SERVER_HOST"
python3 - <<'PY' "$PLAN_FILE"
import json, sys
from pathlib import Path
plan = json.loads(Path(sys.argv[1]).read_text())
print(f"  disk_count={plan['disk_count']}")
print(f"  md_device={plan['md_device']}")
print(f"  fs_type={plan['fs_type']}")
for m in plan['members']:
    print(f"  - {m['disk']}")
PY

run_remote "PAYLOAD_B64=${payload_b64} python3 - <<'PY'
import base64
import json
import os
import shlex
import subprocess

payload = json.loads(base64.b64decode(os.environ['PAYLOAD_B64']).decode())
md_device = payload['md_device']
mount_point = payload['mount_point']
fs_type = payload['fs_type']
members = payload['members']


def run(cmd: str) -> None:
    print('+', cmd, flush=True)
    subprocess.run(cmd, shell=True, check=True)

run('systemctl stop nix-daemon.service >/dev/null 2>&1 || true')
run('systemctl stop nix-daemon.socket >/dev/null 2>&1 || true')
run('umount /mnt/totally-spies-builder/boot-esp-mirror >/dev/null 2>&1 || true')
run('umount /mnt/totally-spies-builder/boot >/dev/null 2>&1 || true')
run('umount /mnt/totally-spies-builder >/dev/null 2>&1 || true')
run('umount /mnt/proc >/dev/null 2>&1 || true')
run('umount /mnt/dev >/dev/null 2>&1 || true')
run('umount /mnt/sys >/dev/null 2>&1 || true')
run('umount /mnt >/dev/null 2>&1 || true')
run('umount /nix >/dev/null 2>&1 || true')
run(f'umount {shlex.quote(mount_point)} >/dev/null 2>&1 || true')
run('mdadm --stop /dev/md127 >/dev/null 2>&1 || true')
run('mdadm --stop --scan >/dev/null 2>&1 || true')
run(f'mdadm --stop {shlex.quote(md_device)} >/dev/null 2>&1 || true')

for member in members:
    disk = member['disk']
    run(f'mdadm --zero-superblock --force {shlex.quote(disk)} >/dev/null 2>&1 || true')
    run(f'wipefs -a -f {shlex.quote(disk)} || true')
    run(f'sgdisk --zap-all {shlex.quote(disk)} >/dev/null 2>&1 || true')

run('udevadm settle')

disk_list = ' '.join(shlex.quote(m['disk']) for m in members)
run(f'mdadm --create {shlex.quote(md_device)} --level=0 --raid-devices={len(members)} {disk_list}')
if fs_type == 'ext4':
    run(f'mkfs.ext4 -F -L builder-root {shlex.quote(md_device)}')
else:
    raise SystemExit(f'Unsupported fs_type: {fs_type}')

run(f'mkdir -p {shlex.quote(mount_point)}')
run(f'mount {shlex.quote(md_device)} {shlex.quote(mount_point)}')
run(f'mdadm --detail {shlex.quote(md_device)}')
run(f'lsblk -o NAME,SIZE,TYPE,FSTYPE,LABEL,MOUNTPOINT {shlex.quote(md_device)} {disk_list}')
PY"