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

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

usage() {
  cat <<'EOF'
Usage: ./ssh.sh <name-or-ip-or-id-prefix> [ssh-args...]

Resolves a running VM by name, guest IP, or ID prefix and opens SSH as root.
EOF
}

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STATE="$DIR/state"
VM_ROOT="$STATE/vms"
SSH_KEY="$DIR/id_ed25519"

find_vm_json() {
  local query="$1"
  local vm_json exact_match="" prefix_match=""
  local exact_count=0 prefix_count=0

  for vm_json in "$VM_ROOT"/*/vm.json; do
    [[ -f "$vm_json" ]] || continue

    local id name guest_ip
    id="$(jq -r '.meta.id // empty' "$vm_json")"
    name="$(jq -r '.meta.name // empty' "$vm_json")"
    guest_ip="$(jq -r '.meta.guest_ip // empty' "$vm_json")"

    if [[ "$guest_ip" == "$query" || "$name" == "$query" || "$id" == "$query" ]]; then
      exact_match="$vm_json"
      exact_count=$((exact_count + 1))
      continue
    fi

    if [[ "$name" == "$query"* || "$id" == "$query"* ]]; then
      prefix_match="$vm_json"
      prefix_count=$((prefix_count + 1))
    fi
  done

  if (( exact_count == 1 )); then
    printf '%s' "$exact_match"
    return 0
  fi
  if (( exact_count > 1 )); then
    log "multiple VMs found for: $query"
    exit 1
  fi
  if (( prefix_count == 1 )); then
    printf '%s' "$prefix_match"
    return 0
  fi
  if (( prefix_count > 1 )); then
    log "multiple VMs found for prefix: $query"
    exit 1
  fi

  log "no VM found for: $query"
  exit 1
}

vm_is_running() {
  local vm_json="$1"
  local pid api_sock
  pid="$(jq -r '.meta.pid // empty' "$vm_json")"
  api_sock="$(jq -r '.meta.api_sock // empty' "$vm_json")"

  [[ -n "$pid" && -n "$api_sock" ]] || return 1
  ps -p "$pid" -o comm=,args= 2>/dev/null | rg -q "firecracker.*--api-sock $api_sock"
}

QUERY="${1:-}"
if [[ "$QUERY" == "-h" || "$QUERY" == "--help" ]]; then
  usage
  exit 0
fi
if [[ -z "$QUERY" ]]; then
  usage
  exit 1
fi
shift

if [[ ! -f "$SSH_KEY" ]]; then
  log "ssh key not found: $SSH_KEY"
  exit 1
fi

VM_JSON="$(find_vm_json "$QUERY")"
if ! vm_is_running "$VM_JSON"; then
  log "vm is not running: $QUERY"
  exit 1
fi

GUEST_IP="$(jq -r '.meta.guest_ip // empty' "$VM_JSON")"
VM_NAME="$(jq -r '.meta.name // empty' "$VM_JSON")"
if [[ -z "$GUEST_IP" ]]; then
  log "guest IP not found for: $QUERY"
  exit 1
fi

log "connecting to $VM_NAME ($GUEST_IP)"
exec ssh \
  -i "$SSH_KEY" \
  -o StrictHostKeyChecking=no \
  -o UserKnownHostsFile=/dev/null \
  "root@$GUEST_IP" \
  "$@"