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

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

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$DIR/dns.sh"
STATE="$DIR/state"
VM_ROOT="$STATE/vms"

usage() {
  cat <<'EOF'
Usage: ./kill.sh <id-or-name-prefix>... [--signal SIGTERM|SIGKILL|...]

Sends a signal to the Firecracker process.
EOF
}

find_vm_json() {
  local query="$1"
  local vm_json match_count=0 match=""

  for vm_json in "$VM_ROOT"/*/vm.json; do
    [[ -f "$vm_json" ]] || continue
    local id name
    id="$(jq -r '.meta.id // empty' "$vm_json")"
    name="$(jq -r '.meta.name // empty' "$vm_json")"
    if [[ "$id" == "$query"* || "$name" == "$query"* ]]; then
      match="$vm_json"
      match_count=$((match_count + 1))
    fi
  done

  if (( match_count == 0 )); then
    log "no VM found for prefix: $query"
    exit 1
  fi
  if (( match_count > 1 )); then
    log "multiple VMs found for prefix: $query"
    exit 1
  fi

  printf '%s' "$match"
}

SIGNAL="TERM"
QUERIES=()
while [[ $# -gt 0 ]]; do
  case "$1" in
    --signal)
      SIGNAL="${2:-}"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      QUERIES+=("$1")
      shift
      ;;
  esac
done

if (( ${#QUERIES[@]} == 0 )); then
  usage
  exit 1
fi

VM_JSONS=()
for query in "${QUERIES[@]}"; do
  VM_JSON="$(find_vm_json "$query")"
  already_added=0
  for existing_vm_json in "${VM_JSONS[@]}"; do
    if [[ "$existing_vm_json" == "$VM_JSON" ]]; then
      already_added=1
      break
    fi
  done
  if (( already_added == 0 )); then
    VM_JSONS+=("$VM_JSON")
  fi
done

for vm_json in "${VM_JSONS[@]}"; do
  vm_id="$(jq -r '.meta.id // empty' "$vm_json")"
  vm_name="$(jq -r '.meta.name // empty' "$vm_json")"
  pid="$(jq -r '.meta.pid // empty' "$vm_json")"
  tap="$(jq -r '.meta.tap // empty' "$vm_json")"
  api_sock="$(jq -r '.meta.api_sock // empty' "$vm_json")"
  base_loop="$(jq -r '.meta.base_loop // empty' "$vm_json")"
  cow_loop="$(jq -r '.meta.cow_loop // empty' "$vm_json")"
  dm_dev="$(jq -r '.meta.dm_dev // empty' "$vm_json")"
  dm_name="$(jq -r '.meta.dm_name // empty' "$vm_json")"
  dns_name="$(jq -r '.meta.dns_name // empty' "$vm_json")"
  if [[ -z "$dns_name" ]]; then
    dns_name="$(banger_dns_name "$vm_name")"
  fi
  if [[ -z "$pid" ]]; then
    log "pid not found in $vm_json"
    exit 1
  fi
  if [[ -z "$api_sock" ]]; then
    log "api_sock not found in $vm_json"
    exit 1
  fi

  if ! ps -p "$pid" -o comm=,args= 2>/dev/null | rg -q "firecracker.*--api-sock $api_sock"; then
    log "pid $pid does not match a running VM for ${vm_name:-$vm_json}"
    exit 1
  fi

  log "sending SIG$SIGNAL to ${vm_name:-$vm_json} (pid $pid)"
  sudo kill "-$SIGNAL" "$pid"
  if ! banger_wait_for_vm_exit "$pid" "$api_sock" 30; then
    log "timed out waiting for ${vm_name:-$vm_json} to exit"
    exit 1
  fi
  banger_teardown_vm_runtime "$tap" "$api_sock" "$dm_name" "$dm_dev" "$cow_loop" "$base_loop"
  banger_mark_vm_stopped "$vm_json"
  banger_dns_remove_record_name "$dns_name"
  log "signal sent to ${vm_name:-$vm_json}"
done