#!/usr/bin/env bash
# Shared state and helpers for the dangerlab clone pool: one frozen base image, N
# single-use overlay clones, one systemd slice per clone.
set -uo pipefail

LAB=/home/user/dangerlab
VAR=/var/lib/buildbox/dangerlab
BASE=$VAR/base.qcow2
SEED=$VAR/seed.iso
RUNDIR=$VAR/run
BASE_SUM=$LAB/base.qcow2.sha256
KEY=$LAB/id_dangerlab
SLOTS=4

# Baked into the frozen base by provision.sh and pinned by
# modules/workstation/claude/tests/lib/danger-guard.sh. Changing either side alone
# makes every exhaustion-class test refuse to run.
LINEAGE_FILE=/etc/dangerlab-lineage
LINEAGE_ID='dangerlab-lineage-v1 f1d0c0de-da4b-4a11-9c3f-0b7e5a2d61a7'

BOOT_ALLOWANCE=420
REAP_GRACE=120
KEEP_TTL=3600
ADMIT_WAIT=3600

V="sudo -n virsh"
SSH="ssh -i $KEY -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -o ConnectTimeout=5"

die()  { printf 'dangerlab: ERROR: %s\n' "$*" >&2; exit 1; }
note() { printf 'dangerlab: %s\n' "$*" >&2; }

slot_dom()      { printf 'dangerlab-slot%s' "$1"; }
slot_dir()      { printf '%s/slot%s' "$VAR" "$1"; }
slot_mac()      { printf '52:54:00:da:1a:0%s' "$1"; }
# Reserved in libvirt's default network by provision.sh. MAC, address and slot are one
# fixed bijection, so an address can never resolve to another agent's clone — which a
# lease lookup could, once every clone mints its own DHCP client identifier.
slot_ip()       { printf '192.168.122.10%s' "$1"; }
slot_lock()     { printf '%s/slot%s.lock' "$RUNDIR" "$1"; }
slot_deadline() { printf '%s/slot%s.deadline' "$RUNDIR" "$1"; }
slot_keep()     { printf '%s/slot%s.keep' "$RUNDIR" "$1"; }
slot_overlay()  { printf '%s/slot%s/overlay.qcow2' "$VAR" "$1"; }
slot_console()  { printf '%s/slot%s/console.log' "$VAR" "$1"; }

slot_seq() { seq 1 "$SLOTS"; }

# The base is the only thing a clone inherits. An unverified base produces no clones.
base_verify() {
  [[ -f "$BASE" ]] || die "base image missing: $BASE"
  [[ -f "$BASE_SUM" ]] || die "base checksum pin missing: $BASE_SUM"
  sudo -n sha256sum -c "$BASE_SUM" >/dev/null 2>&1 \
    || die "base image checksum MISMATCH against $BASE_SUM — refusing to clone"
}

vm_state() { $V domstate "$1" 2>/dev/null | head -1; }

domain_kill() {
  local dom="$1" s
  s=$(vm_state "$dom")
  if [[ -n "$s" && "$s" != "shut off" ]]; then
    $V destroy "$dom" >/dev/null 2>&1
    for _ in $(seq 1 30); do
      [[ "$(vm_state "$dom")" == "shut off" || -z "$(vm_state "$dom")" ]] && break
      sleep 1
    done
  fi
  [[ -z "$(vm_state "$dom")" ]] || $V undefine "$dom" --nvram >/dev/null 2>&1 || $V undefine "$dom" >/dev/null 2>&1
  [[ -z "$(vm_state "$dom")" ]]
}

gen_domain() { # slot
  local k="$1"
  cat <<XML
<domain type='kvm'>
  <name>$(slot_dom "$k")</name>
  <memory unit='KiB'>2097152</memory>
  <currentMemory unit='KiB'>2097152</currentMemory>
  <memtune>
    <hard_limit unit='KiB'>2883584</hard_limit>
    <swap_hard_limit unit='KiB'>2883584</swap_hard_limit>
  </memtune>
  <vcpu placement='static'>2</vcpu>
  <resource><partition>/machine/dangerlab/slot$k</partition></resource>
  <cputune>
    <global_period>100000</global_period>
    <global_quota>200000</global_quota>
    <shares>256</shares>
  </cputune>
  <os>
    <type arch='x86_64' machine='q35'>hvm</type>
    <boot dev='hd'/>
  </os>
  <features><acpi/><apic/></features>
  <cpu mode='host-passthrough' check='none'/>
  <clock offset='utc'/>
  <on_poweroff>destroy</on_poweroff>
  <on_reboot>restart</on_reboot>
  <on_crash>destroy</on_crash>
  <devices>
    <emulator>/usr/bin/qemu-system-x86_64</emulator>
    <disk type='file' device='disk'>
      <driver name='qemu' type='qcow2'/>
      <source file='$(slot_overlay "$k")'/>
      <target dev='vda' bus='virtio'/>
    </disk>
    <disk type='file' device='cdrom'>
      <driver name='qemu' type='raw'/>
      <source file='$SEED'/>
      <target dev='sda' bus='sata'/>
      <readonly/>
    </disk>
    <interface type='network'>
      <source network='default'/>
      <mac address='$(slot_mac "$k")'/>
      <model type='virtio'/>
    </interface>
    <serial type='file'>
      <source path='$(slot_console "$k")'/>
      <target port='0'/>
    </serial>
    <console type='file'>
      <source path='$(slot_console "$k")'/>
      <target type='serial' port='0'/>
    </console>
    <channel type='unix'>
      <target type='virtio' name='org.qemu.guest_agent.0'/>
    </channel>
    <graphics type='vnc' port='-1' listen='127.0.0.1'/>
    <video><model type='vga'/></video>
    <memballoon model='none'/>
  </devices>
</domain>
XML
}

# Destroy the domain and discard the overlay. Safe to call on a slot that holds
# nothing. The run markers are the owner's and the reaper's to clear, never this
# function's: clone_create calls it, and clearing them here erased the deadline the
# reaper judges a live clone by.
clone_destroy() { # slot
  local k="$1"
  domain_kill "$(slot_dom "$k")" || return 1
  sudo -n rm -f "$(slot_overlay "$k")"
  return 0
}

# Build a single-use clone from the frozen base. Every precondition is read back from
# the host; nothing is assumed.
clone_create() { # slot
  local k="$1" dir dom xml
  dir=$(slot_dir "$k"); dom=$(slot_dom "$k")
  clone_destroy "$k" || die "slot$k: could not clear the previous clone"
  sudo -n install -d -o libvirt-qemu -g libvirt-qemu -m 0755 "$dir"
  sudo -n qemu-img create -f qcow2 -F qcow2 -b "$BASE" "$(slot_overlay "$k")" >/dev/null \
    || die "slot$k: overlay create failed"
  sudo -n chown libvirt-qemu:libvirt-qemu "$(slot_overlay "$k")"
  sudo -n qemu-img info --backing-chain "$(slot_overlay "$k")" 2>/dev/null \
    | grep -qF "backing file: $BASE" || die "slot$k: overlay is not backed by $BASE — refusing"
  sudo -n truncate -s 0 "$(slot_console "$k")" 2>/dev/null \
    || { sudo -n touch "$(slot_console "$k")"; sudo -n chown libvirt-qemu:libvirt-qemu "$(slot_console "$k")"; }
  xml=$(mktemp)
  gen_domain "$k" >"$xml"
  $V define "$xml" >/dev/null || { rm -f "$xml"; die "slot$k: virsh define failed"; }
  rm -f "$xml"
  $V start "$dom" >/dev/null 2>&1 || die "slot$k: virsh start failed"
}

# ---------------------------------------------------------------------------
# Slot API. Any host-side script that needs a disposable guest uses this; the
# slot is released by the shell's own exit, so a crashed caller frees it and the
# reaper collects what it left behind.
#
#   source /home/user/dangerlab/dangerlab-lib.sh
#   lab_acquire 600          # admission + a booted clone; sets LAB_SLOT, LAB_IP
#   lab_ssh 'uname -a'
#   lab_reset                # discard this clone, boot another in the same slot
#   lab_release              # optional; the EXIT trap does it anyway
# ---------------------------------------------------------------------------
LAB_SLOT=""; LAB_IP=""; LAB_LOCKFD=""; LAB_BUDGET=300; LAB_KEEP=0

lab_acquire() { # [budget-seconds]
  local waited=0 k fd
  LAB_BUDGET="${1:-300}"
  [[ -d "$RUNDIR" && -w "$RUNDIR" ]] || die "$RUNDIR missing or unwritable — run provision.sh"
  base_verify
  while :; do
    for k in $(slot_seq); do
      exec {fd}>"$(slot_lock "$k")" || continue
      if flock -n "$fd"; then LAB_SLOT="$k"; LAB_LOCKFD="$fd"; break 2; fi
      exec {fd}>&-
    done
    (( waited == 0 )) && note "all $SLOTS clone slots busy — waiting for one (the cap is host capacity)"
    (( waited >= ADMIT_WAIT )) && die "no clone slot free after ${ADMIT_WAIT}s"
    sleep 5; waited=$(( waited + 5 ))
  done
  trap _lab_exit EXIT
  trap 'exit 130' INT
  trap 'exit 143' TERM
  lab_reset
}

lab_reset() {
  [[ -n "$LAB_SLOT" ]] || die "lab_reset called before lab_acquire"
  # Written before the clone exists, so a live clone always carries a deadline the
  # reaper can judge it by.
  printf '%s\n' "$(( $(date +%s) + BOOT_ALLOWANCE + LAB_BUDGET ))" >"$(slot_deadline "$LAB_SLOT")"
  rm -f "$(slot_keep "$LAB_SLOT")"
  clone_create "$LAB_SLOT"
  LAB_IP=$(wait_ssh "$LAB_SLOT" "$BOOT_ALLOWANCE") \
    || die "slot$LAB_SLOT: clone did not boot within ${BOOT_ALLOWANCE}s"
}

lab_ssh()  { $SSH lab@"$LAB_IP" "$@"; }
lab_scp()  { scp -q -i "$KEY" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
                 -o LogLevel=ERROR -r "$1" lab@"$LAB_IP":"$2"; }
lab_wait_ssh() { LAB_IP=$(wait_ssh "$LAB_SLOT" "${1:-$BOOT_ALLOWANCE}"); }

lab_release() {
  [[ -n "$LAB_SLOT" ]] || return 0
  if [[ $LAB_KEEP -eq 1 ]]; then
    printf '%s\n' "$(( $(date +%s) + KEEP_TTL ))" >"$(slot_keep "$LAB_SLOT")"
    note "slot$LAB_SLOT pinned for inspection; the reaper collects it in ${KEEP_TTL}s"
  else
    clone_destroy "$LAB_SLOT" \
      || printf 'dangerlab: WARNING: slot%s not fully cleared; the reaper will collect it\n' "$LAB_SLOT" >&2
  fi
  [[ $LAB_KEEP -eq 1 ]] || rm -f "$(slot_keep "$LAB_SLOT")"
  rm -f "$(slot_deadline "$LAB_SLOT")"
  LAB_SLOT=""; LAB_IP=""
}

_lab_exit() { local rc=$?; trap - EXIT; lab_release; exit "$rc"; }

wait_ssh() { # slot seconds -> prints ip
  local k="$1" deadline=$(( SECONDS + ${2:-300} )) ip
  ip=$(slot_ip "$k")
  while (( SECONDS < deadline )); do
    if $SSH lab@"$ip" true 2>/dev/null; then printf '%s' "$ip"; return 0; fi
    sleep 3
  done
  return 1
}

# Declares the MAC -> address reservations the pool depends on. Idempotent; only ever
# touches the entries whose MACs belong to this lab.
net_reserve() {
  local k mac ip
  for k in $(slot_seq); do
    mac=$(slot_mac "$k"); ip=$(slot_ip "$k")
    $V net-update default delete ip-dhcp-host "<host mac='$mac'/>" --live --config >/dev/null 2>&1 || true
    $V net-update default add ip-dhcp-host "<host mac='$mac' name='dangerlab-slot$k' ip='$ip'/>" --live --config >/dev/null \
      || die "could not reserve $ip for $mac in libvirt's default network"
  done
}