#!/usr/bin/env bash
set -euo pipefail

log() {
  printf '[spawn] %s\n' "$*"
}

usage() {
  cat <<'EOF'
Usage: ./run.sh [options]

Options:
  --name <name>        VM name (lowercase letters, digits, -, _)
  --vcpu <count>       vCPU count (default: 2)
  --ram <mib>          RAM in MiB (default: 1024)
  --disk-size <size>   Root disk size (e.g. 4G, 10240M); must be >= base image
  -h, --help           Show this help
EOF
}

log "starting"

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STATE="$DIR/state"
mkdir -p "$STATE"

FC_BIN="$DIR/firecracker"
KERNEL="$DIR/vmlinux"
ROOTFS="$DIR/rootfs.ext4"
SSH_KEY="$DIR/id_ed25519"
NAMEGEN="$DIR/namegen"

BR_DEV="br-fc"
BR_IP="172.16.0.1"
CIDR="24"

DEFAULT_VCPU=2
DEFAULT_RAM=1024
MAX_VCPU=16
MAX_RAM=32768
MAX_DISK_BYTES=$((128 * 1024 * 1024 * 1024))

VCPU_COUNT="$DEFAULT_VCPU"
RAM_MIB="$DEFAULT_RAM"
DISK_SIZE=""
VM_NAME=""

shopt -s nullglob

name_taken() {
  local candidate="$1"
  local info existing_name
  for info in "$STATE"/vm-*/info; do
    existing_name="$(awk -F= '$1=="name"{print $2}' "$info")"
    if [[ "$existing_name" == "$candidate" ]]; then
      return 0
    fi
  done
  return 1
}

parse_disk_size() {
  local raw="$1"
  if [[ "$raw" =~ ^([0-9]+)([KMG])?$ ]]; then
    local num="${BASH_REMATCH[1]}"
    local unit="${BASH_REMATCH[2]}"
    case "$unit" in
      K) echo $((num * 1024)) ;;
      M|"") echo $((num * 1024 * 1024)) ;;
      G) echo $((num * 1024 * 1024 * 1024)) ;;
    esac
    return 0
  fi
  return 1
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --name)
      VM_NAME="${2:-}"
      shift 2
      ;;
    --vcpu)
      VCPU_COUNT="${2:-}"
      shift 2
      ;;
    --ram)
      RAM_MIB="${2:-}"
      shift 2
      ;;
    --disk-size)
      DISK_SIZE="${2:-}"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      log "unknown option: $1"
      usage
      exit 1
      ;;
  esac
done

if ! [[ "$VCPU_COUNT" =~ ^[0-9]+$ ]]; then
  log "invalid --vcpu value: $VCPU_COUNT"
  exit 1
fi
if (( VCPU_COUNT < 1 || VCPU_COUNT > MAX_VCPU )); then
  log "vcpu must be between 1 and $MAX_VCPU"
  exit 1
fi

if ! [[ "$RAM_MIB" =~ ^[0-9]+$ ]]; then
  log "invalid --ram value: $RAM_MIB"
  exit 1
fi
if (( RAM_MIB < 256 || RAM_MIB > MAX_RAM )); then
  log "ram must be between 256 and $MAX_RAM MiB"
  exit 1
fi

DISK_BYTES=""
if [[ -n "$DISK_SIZE" ]]; then
  if ! DISK_BYTES="$(parse_disk_size "$DISK_SIZE")"; then
    log "invalid --disk-size value: $DISK_SIZE"
    exit 1
  fi
  if (( DISK_BYTES > MAX_DISK_BYTES )); then
    log "disk-size exceeds max of $((MAX_DISK_BYTES / 1024 / 1024 / 1024))G"
    exit 1
  fi
fi

if [[ -z "$VM_NAME" ]]; then
  if [[ ! -x "$NAMEGEN" ]]; then
    log "name generator not found: $NAMEGEN"
    exit 1
  fi
  for _ in $(seq 1 20); do
    VM_NAME="$("$NAMEGEN")"
    if ! name_taken "$VM_NAME"; then
      break
    fi
  done
fi

if [[ -z "$VM_NAME" ]]; then
  log "failed to generate a unique name"
  exit 1
fi

if ! [[ "$VM_NAME" =~ ^[a-z0-9][a-z0-9_-]{0,63}$ ]]; then
  log "invalid --name value: $VM_NAME"
  exit 1
fi
if name_taken "$VM_NAME"; then
  log "name already in use: $VM_NAME"
  exit 1
fi

VM_ID="$(head -c 32 /dev/urandom | xxd -p -c 256)"
VM_TAG="${VM_ID:0:8}"
VM_DIR="$STATE/vm-$VM_ID"
mkdir -p "$VM_DIR"

API_SOCK="$STATE/fc-$VM_TAG.sock"
LOG_FILE="$VM_DIR/firecracker.log"
TAP_DEV="tap-fc-$VM_TAG"

# Allocate guest IP
NEXT_IP_FILE="$STATE/next_ip"
NEXT_IP="$(cat "$NEXT_IP_FILE" 2>/dev/null || echo 2)"
GUEST_IP="172.16.0.$NEXT_IP"
echo "$((NEXT_IP + 1))" > "$NEXT_IP_FILE"

log "vm id: $VM_ID"
log "vm name: $VM_NAME"
log "allocated guest ip: $GUEST_IP"

sudo -v

VM_STARTED=0
CLEANUP_ON_EXIT=0
KEEP_VM_DIR_ON_FAIL=1
DISK_PATH="$ROOTFS"

cleanup() {
  local exit_code=$?
  if [[ "$VM_STARTED" -eq 1 && "$CLEANUP_ON_EXIT" -eq 0 ]]; then
    return
  fi
  log "cleaning up"
  if [[ -n "${FC_PID:-}" ]]; then
    sudo kill "$FC_PID" 2>/dev/null || true
  fi
  if [[ -n "${TAP_DEV:-}" ]]; then
    sudo ip link del "$TAP_DEV" 2>/dev/null || true
  fi
  rm -f "${API_SOCK:-}"
  if [[ -n "${VM_DIR:-}" && -d "$VM_DIR" ]]; then
    if [[ "$KEEP_VM_DIR_ON_FAIL" -eq 0 || "$VM_STARTED" -eq 1 || "$CLEANUP_ON_EXIT" -eq 1 ]]; then
      rm -rf "$VM_DIR"
    fi
  fi
  return "$exit_code"
}

on_signal() {
  CLEANUP_ON_EXIT=1
  exit 1
}

trap cleanup EXIT
trap on_signal INT TERM

FC_USE_SUDO="${FC_USE_SUDO:-1}"
FC_RUN=("$FC_BIN")
CURL_CMD=(curl)
if [[ "$FC_USE_SUDO" == "1" ]]; then
  log "running firecracker with sudo (FC_USE_SUDO=1)"
  FC_RUN=(sudo -E "$FC_BIN")
  CURL_CMD=(sudo -E curl)
fi

if command -v setcap >/dev/null 2>&1; then
  if ! getcap "$FC_BIN" 2>/dev/null | rg -q "cap_net_admin"; then
    log "granting cap_net_admin to firecracker binary"
    sudo setcap cap_net_admin+ep "$FC_BIN"
  fi
else
  log "setcap not available; firecracker may need root to open TAP"
fi

if [[ -n "$DISK_BYTES" ]]; then
  if ! command -v resize2fs >/dev/null 2>&1; then
    log "resize2fs required for --disk-size"
    exit 1
  fi
  DISK_PATH="$VM_DIR/rootfs.ext4"
  cp --reflink=auto "$ROOTFS" "$DISK_PATH"
  BASE_BYTES="$(stat -c%s "$ROOTFS")"
  if (( DISK_BYTES < BASE_BYTES )); then
    log "disk-size must be >= base image size"
    exit 1
  fi
  if (( DISK_BYTES > BASE_BYTES )); then
    log "resizing rootfs to $DISK_SIZE"
    truncate -s "$DISK_BYTES" "$DISK_PATH"
    resize2fs "$DISK_PATH" >/dev/null
  fi
fi

# Host bridge
if ! ip link show "$BR_DEV" >/dev/null 2>&1; then
  log "creating host bridge $BR_DEV ($BR_IP/$CIDR)"
  sudo ip link add name "$BR_DEV" type bridge
  sudo ip addr add "${BR_IP}/${CIDR}" dev "$BR_DEV"
  sudo ip link set "$BR_DEV" up
else
  log "host bridge $BR_DEV already exists"
  # Ensure existing bridge is up in case it was left down.
  sudo ip link set "$BR_DEV" up
fi

# Per-VM TAP
log "creating tap device $TAP_DEV"
TAP_USER="${SUDO_UID:-$(id -u)}"
TAP_GROUP="${SUDO_GID:-$(id -g)}"
sudo ip tuntap add dev "$TAP_DEV" mode tap user "$TAP_USER" group "$TAP_GROUP"
sudo ip link set "$TAP_DEV" master "$BR_DEV"
sudo ip link set "$TAP_DEV" up
sudo ip link set "$BR_DEV" up

# Start Firecracker
log "starting firecracker process"
rm -f "$API_SOCK"
nohup "${FC_RUN[@]}" --api-sock "$API_SOCK" >"$LOG_FILE" 2>&1 &
FC_PID="$!"
log "firecracker pid: $FC_PID"

# Wait for API socket
log "waiting for firecracker api socket"
for _ in $(seq 1 200); do
  [[ -S "$API_SOCK" ]] && break
  sleep 0.02
done
if [[ ! -S "$API_SOCK" ]]; then
  log "firecracker api socket not ready"
  if [[ -f "$LOG_FILE" ]]; then
    log "firecracker log (tail)"
    tail -n 20 "$LOG_FILE" || true
  fi
  exit 1
fi
log "api socket ready"

if [[ "$FC_USE_SUDO" == "1" ]]; then
  SUDO_CHILD_PID="$(pgrep -n -f "$API_SOCK" || true)"
  if [[ -n "$SUDO_CHILD_PID" ]]; then
    FC_PID="$SUDO_CHILD_PID"
    log "firecracker child pid: $FC_PID"
  fi
fi
echo "$FC_PID" > "$VM_DIR/pid"

# Machine config
log "configuring machine"
"${CURL_CMD[@]}" --unix-socket "$API_SOCK" -X PUT http://localhost/machine-config \
  -H "Content-Type: application/json" \
  -d "{
    \"vcpu_count\": $VCPU_COUNT,
    \"mem_size_mib\": $RAM_MIB,
    \"smt\": false
  }" >/dev/null

# Boot source
log "configuring boot source"
KCMD="console=ttyS0 reboot=k panic=1 pci=off root=/dev/vda rw ip=${GUEST_IP}::${BR_IP}:255.255.255.0::eth0:off"

"${CURL_CMD[@]}" --unix-socket "$API_SOCK" -X PUT http://localhost/boot-source \
  -H "Content-Type: application/json" \
  -d "{
    \"kernel_image_path\": \"$KERNEL\",
    \"boot_args\": \"$KCMD\"
  }" >/dev/null

# Root filesystem
log "attaching root filesystem"
"${CURL_CMD[@]}" --unix-socket "$API_SOCK" -X PUT http://localhost/drives/rootfs \
  -H "Content-Type: application/json" \
  -d "{
    \"drive_id\": \"rootfs\",
    \"path_on_host\": \"$DISK_PATH\",
    \"is_root_device\": true,
    \"is_read_only\": false
  }" >/dev/null

# Network interface
log "configuring network interface"
"${CURL_CMD[@]}" --unix-socket "$API_SOCK" -X PUT http://localhost/network-interfaces/eth0 \
  -H "Content-Type: application/json" \
  -d "{
    \"iface_id\": \"eth0\",
    \"host_dev_name\": \"$TAP_DEV\"
  }" >/dev/null

# Start VM
log "starting virtual machine"
"${CURL_CMD[@]}" --unix-socket "$API_SOCK" -X PUT http://localhost/actions \
  -H "Content-Type: application/json" \
  -d '{ "action_type": "InstanceStart" }' >/dev/null
VM_STARTED=1

cat > "$VM_DIR/info" <<EOF
id=$VM_ID
name=$VM_NAME
pid=$FC_PID
created_at=$(date -Iseconds)
guest_ip=$GUEST_IP
tap=$TAP_DEV
api_sock=$API_SOCK
log=$LOG_FILE
EOF

log "vm started successfully"
log "guest ip: $GUEST_IP"
log "ssh: ssh -i \"$SSH_KEY\" root@$GUEST_IP"
log "logs: $LOG_FILE"