#!/bin/bash
# microvm — Firecracker-class microVMs on ultix-streaming (Redmine #615)
# Engine: qemu -M microvm + KVM (default) | Firecracker (MICROVM_ENGINE=firecracker;
# needs a KVM exposing USER_MEMORY/SET_TSS_ADDR/VAPIC caps — refused by this
# host's nested KVM until the #601 hardware upgrade; binary stays installed).
# Guests: Debian 13 cloud image, cloud-init NoCloud seed, static tap networking.
# Instances: ~/.microvm/<name>. Host prep: microvm-net.service (taps mv-<p><i>).

set -euo pipefail
export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

BASE=/usr/local/lib/microvm/base
ROOT=$HOME/.microvm
DISK_GIB=${DISK_GIB:-10}
SSH_USER=debian

case "$USER" in
  reachableceo)           PREFIX=r; SUBNET=172.30.0 ;;
  reachableceo-offstage)  PREFIX=o; SUBNET=172.30.1 ;;
  *) echo "microvm: no tap pool for user '$USER' (add one in microvm-net.sh)" >&2; exit 1 ;;
esac

die() { echo "microvm: $*" >&2; exit 1; }

slot_ip() { # slot -> "host_ip guest_ip"
  echo "$SUBNET.$((4*$1+1))" "$SUBNET.$((4*$1+2))"
}

free_slot() {
  local used slot
  used=$(cat "$ROOT"/*/vm.conf 2>/dev/null | grep -E '^SLOT=' | cut -d= -f2 | sort -u || true)
  for slot in 0 1 2 3; do
    grep -qx "$slot" <<<"$used" || { echo "$slot"; return; }
  done
  die "all 4 tap slots busy; remove an instance first (microvm list / rm)"
}

load_conf() {
  [ -f "$ROOT/$1/vm.conf" ] || die "no such instance '$1' (microvm list)"
  . "$ROOT/$1/vm.conf"
}

cmd_create() {
  local name=${1:-} mem=${2:-1024} vcpu=${3:-1} slot hostip guestip pubkey
  [[ "$name" =~ ^[a-z][a-z0-9-]{0,15}$ ]] || die "name: ^[a-z][a-z0-9-]{0,15}$"
  [ -d "$ROOT/$name" ] && die "instance '$name' already exists"
  slot=$(free_slot)
  read -r hostip guestip <<<"$(slot_ip "$slot")"
  mkdir -p "$ROOT/$name"

  pubkey=$(ls -1 "$HOME"/.ssh/id_*.pub 2>/dev/null | head -1)
  [ -n "$pubkey" ] || die "no $HOME/.ssh/id_*.pub found (needed for guest login)"

  cat > "$ROOT/$name/user-data" <<EOF
#cloud-config
hostname: $name
users:
  - name: $SSH_USER
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/bash
    lock_passwd: true
    ssh_authorized_keys:
      - $(cat "$pubkey")
EOF
  cat > "$ROOT/$name/network-config" <<EOF
version: 2
ethernets:
  eth0:
    addresses: [$guestip/30]
    routes: [{via: $hostip, to: default}]
    nameservers: {addresses: [1.1.1.1, 9.9.9.9]}
EOF
  printf 'instance-id: microvm-%s-%s\nlocal-hostname: %s\n' "$name" "$(date +%s)" "$name" > "$ROOT/$name/meta-data"

  truncate -s 4M "$ROOT/$name/seed.img"
  mkfs.vfat -n CIDATA "$ROOT/$name/seed.img" >/dev/null
  mcopy -i "$ROOT/$name/seed.img" "$ROOT/$name/user-data" "$ROOT/$name/network-config" "$ROOT/$name/meta-data" ::

  echo "microvm: cloning base rootfs (${DISK_GIB}G sparse)..."
  qemu-img convert -O raw "$BASE/base-rootfs.raw" "$ROOT/$name/rootfs.raw"
  qemu-img resize "$ROOT/$name/rootfs.raw" "${DISK_GIB}G" >/dev/null

  if [ "$PREFIX" = r ]; then mnet=0; else mnet=1; fi
  MAC=$(printf '52:54:00:30:%02x:%02x' $((mnet*16+slot)) $((slot+1)))

  cat > "$ROOT/$name/vm.conf" <<EOF
NAME=$name
SLOT=$slot
MEM=$mem
VCPU=$vcpu
GUEST_IP=$guestip
HOST_IP=$hostip
TAP=mv-$PREFIX$slot
MAC=$MAC
ENGINE=\${MICROVM_ENGINE:-qemu}
EOF
  echo "microvm: '$name' created — guest $guestip (gw $hostip, tap mv-$PREFIX$slot), $mem MiB / $vcpu vCPU"
  echo "microvm: start it:   microvm start $name && microvm ssh $name"
}

start_qemu() {
  nohup qemu-system-x86_64 \
    -accel kvm -cpu host -M microvm \
    -m "$MEM" -smp "$VCPU" \
    -kernel "$BASE/vmlinuz" -initrd "$BASE/initrd" \
    -append "console=ttyS0 reboot=k panic=-1 i8042.noaux random.trust_cpu=on root=/dev/vda1 rw" \
    -drive file="$ROOT/$NAME/rootfs.raw",if=none,id=r0,format=raw \
    -device virtio-blk-device,drive=r0 \
    -drive file="$ROOT/$NAME/seed.img",if=none,id=s0,format=raw,readonly=on \
    -device virtio-blk-device,drive=s0 \
    -netdev tap,id=n0,ifname="$TAP",script=no,downscript=no \
    -device virtio-net-device,netdev=n0,mac="$MAC" \
    -serial "file:$ROOT/$NAME/serial.log" \
    -display none -no-reboot \
    >"$ROOT/$NAME/qemu.log" 2>&1 &
  echo $! > "$ROOT/$NAME/pid"
}

start_fc() {
  cat > "$ROOT/$NAME/config.json" <<EOF
{
  "boot-source": {
    "kernel_image_path": "$BASE/vmlinuz",
    "initrd_path": "$BASE/initrd",
    "boot_args": "console=ttyS0 reboot=k panic=-1 i8042.noaux random.trust_cpu=on root=/dev/vda1 rw"
  },
  "drives": [
    {"drive_id": "rootfs", "path_on_host": "$ROOT/$NAME/rootfs.raw", "is_root_device": true, "is_read_only": false},
    {"drive_id": "seed", "path_on_host": "$ROOT/$NAME/seed.img", "is_root_device": false, "is_read_only": true}
  ],
  "network-interfaces": [{"iface_id": "eth0", "guest_mac": "$MAC", "host_dev_name": "$TAP"}],
  "machine-config": {"vcpu_count": $VCPU, "mem_size_mib": $MEM, "smt": false},
  "serial": {"serial_out_path": "$ROOT/$NAME/serial.log"}
}
EOF
  nohup firecracker --api-sock "$ROOT/$NAME/fc.sock" --config-file "$ROOT/$NAME/config.json" \
    >"$ROOT/$NAME/fc.log" 2>&1 &
  echo $! > "$ROOT/$NAME/pid"
}

running() { [ -f "$ROOT/$1/pid" ] && kill -0 "$(cat "$ROOT/$1/pid")" 2>/dev/null; }

cmd_start() {
  load_conf "$1"
  running "$NAME" && die "'$NAME' already running (pid $(cat "$ROOT/$NAME/pid"))"
  ip link show "$TAP" >/dev/null 2>&1 || die "tap $TAP missing — run: systemctl start microvm-net"
  rm -f "$ROOT/$NAME/serial.log"
  if [ "$ENGINE" != firecracker ] && ! id -nG | tr ' ' '\n' | grep -qx kvm; then
    exec sg kvm -c "$0 start $NAME"
  fi
  if [ "$ENGINE" = firecracker ]; then start_fc; else start_qemu; fi
  echo "microvm: '$NAME' starting — console: microvm console $NAME; ssh: microvm ssh $NAME (cloud-init ~10-20s)"
}

cmd_ssh() {
  load_conf "$1"; shift
  exec ssh -o UserKnownHostsFile="$ROOT/$NAME/known_hosts" -o StrictHostKeyChecking=accept-new \
    -o ConnectTimeout=5 "$SSH_USER@$GUEST_IP" "$@"
}

cmd_console() {
  load_conf "$1"
  [ -f "$ROOT/$NAME/serial.log" ] || die "no serial.log yet (start first)"
  exec tail -f "$ROOT/$NAME/serial.log"
}

cmd_stop() {
  load_conf "$1"
  running "$NAME" || { echo "microvm: '$NAME' not running"; return 0; }
  ssh -o UserKnownHostsFile="$ROOT/$NAME/known_hosts" -o ConnectTimeout=3 -o BatchMode=yes \
    "$SSH_USER@$GUEST_IP" "sudo poweroff" 2>/dev/null || true
  for _ in $(seq 1 20); do
    running "$NAME" || { echo "microvm: '$NAME' stopped"; return 0; }
    sleep 1
  done
  kill "$(cat "$ROOT/$NAME/pid")" 2>/dev/null || true
  echo "microvm: '$NAME' force-stopped"
}

cmd_list() {
  printf '%-16s %-8s %-15s %-9s %6s %5s  %s\n' NAME STATE IP TAP MEM VCPUT ENGINE
  for conf in "$ROOT"/*/vm.conf; do
    [ -f "$conf" ] || continue
    (. "$conf"
     if running "$NAME"; then st=running; else st=stopped; fi
     printf '%-16s %-8s %-15s %-9s %6s %5s  %s\n' "$NAME" "$st" "$GUEST_IP" "$TAP" "$MEM" "$VCPU" "$ENGINE")
  done
}

cmd_rm() {
  load_conf "$1"
  running "$NAME" && die "'$NAME' is running — stop it first"
  rm -rf "$ROOT/$NAME"
  echo "microvm: '$NAME' removed"
}

case "${1:-}" in
  create) shift; [ $# -ge 1 ] || die "usage: microvm create NAME [MEM_MIB] [VCPU]"; cmd_create "$@" ;;
  start)  shift; [ $# -ge 1 ] || die "usage: microvm start NAME"; cmd_start "$@" ;;
  ssh)    shift; [ $# -ge 1 ] || die "usage: microvm ssh NAME [cmd...]"; cmd_ssh "$@" ;;
  console) shift; [ $# -ge 1 ] || die "usage: microvm console NAME"; cmd_console "$@" ;;
  stop)   shift; [ $# -ge 1 ] || die "usage: microvm stop NAME"; cmd_stop "$@" ;;
  list)   cmd_list ;;
  rm)     shift; [ $# -ge 1 ] || die "usage: microvm rm NAME"; cmd_rm "$@" ;;
  *) sed -n '2,8p' "$0"; echo; echo "usage: microvm {create|start|ssh|console|stop|list|rm} NAME [args]"; exit 1 ;;
esac
