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

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

usage() {
  cat <<'EOF'
Usage: ./restore.sh <id-or-name-prefix>

Restarts a VM using existing disks and COW snapshot.
EOF
}

find_vm_meta() {
  local query="$1"
  local meta_json match_count=0 match=""

  for meta_json in state/vms/*/meta.json; do
    [[ -f "$meta_json" ]] || continue
    local id name
    id="$(jq -r '.id // empty' "$meta_json")"
    name="$(jq -r '.name // empty' "$meta_json")"
    if [[ "$id" == "$query"* || "$name" == "$query"* ]]; then
      match="$meta_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"
}

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

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FC_BIN="$DIR/firecracker"

if ! command -v jq >/dev/null 2>&1; then
  log "jq is required"
  exit 1
fi
if ! command -v dmsetup >/dev/null 2>&1 || ! command -v losetup >/dev/null 2>&1 || ! command -v blockdev >/dev/null 2>&1; then
  log "dmsetup, losetup, and blockdev are required"
  exit 1
fi
if ! command -v e2cp >/dev/null 2>&1 || ! command -v e2rm >/dev/null 2>&1; then
  log "e2cp and e2rm are required"
  exit 1
fi

META_JSON="$(find_vm_meta "$QUERY")"
VM_DIR="$(dirname "$META_JSON")"
VM_JSON="$VM_DIR/vm.json"

ROOTFS="$(jq -r '.rootfs // empty' "$META_JSON")"
KERNEL="$(jq -r '.kernel // empty' "$META_JSON")"
HOME_PATH="$(jq -r '.home_path // empty' "$META_JSON")"
VAR_PATH="$(jq -r '.var_path // empty' "$META_JSON")"
TAP_DEV="$(jq -r '.tap // empty' "$META_JSON")"
API_SOCK="$(jq -r '.api_sock // empty' "$META_JSON")"
LOG_FILE="$(jq -r '.log // empty' "$META_JSON")"
GUEST_IP="$(jq -r '.guest_ip // empty' "$META_JSON")"
DM_NAME="$(jq -r '.dm_name // empty' "$META_JSON")"
COW_FILE="$(jq -r '.cow_file // empty' "$META_JSON")"

if [[ -z "$ROOTFS" || -z "$KERNEL" || -z "$HOME_PATH" || -z "$VAR_PATH" || -z "$API_SOCK" || -z "$TAP_DEV" || -z "$GUEST_IP" || -z "$DM_NAME" || -z "$COW_FILE" ]]; then
  log "meta.json missing required fields"
  exit 1
fi
if [[ ! -f "$ROOTFS" || ! -f "$KERNEL" || ! -f "$HOME_PATH" || ! -f "$VAR_PATH" || ! -f "$COW_FILE" || ! -f "$FC_BIN" ]]; then
  log "missing disk/kernel file(s)"
  exit 1
fi
if [[ ! -f "$VM_JSON" ]]; then
  log "vm.json missing: $VM_JSON"
  exit 1
fi

sudo -v

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

VM_STARTED=0
cleanup() {
  if [[ "$VM_STARTED" -eq 1 ]]; then
    return
  fi
  sudo ip link del "$TAP_DEV" 2>/dev/null || true
  rm -f "$API_SOCK"
}
trap cleanup EXIT

# 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
  sudo ip link set "$BR_DEV" up
fi

sock_dir="$(dirname "$API_SOCK")"
sudo mkdir -p "$sock_dir"
sudo chown "$(id -u):$(id -g)" "$sock_dir"

# Recreate dm-snapshot
sudo dmsetup remove "$DM_NAME" 2>/dev/null || true
BASE_LOOP="$(sudo losetup -f --show --read-only "$ROOTFS")"
COW_LOOP="$(sudo losetup -f --show "$COW_FILE")"
SECTORS="$(sudo blockdev --getsz "$BASE_LOOP")"
sudo dmsetup create "$DM_NAME" --table "0 $SECTORS snapshot $BASE_LOOP $COW_LOOP P 8"
DM_DEV="/dev/mapper/$DM_NAME"

jq \
  --arg base_loop "$BASE_LOOP" \
  --arg cow_loop "$COW_LOOP" \
  --arg dm_dev "$DM_DEV" \
  '.base_loop=$base_loop | .cow_loop=$cow_loop | .dm_dev=$dm_dev' \
  "$META_JSON" > "$META_JSON.tmp" && mv "$META_JSON.tmp" "$META_JSON"

# Update /etc/resolv.conf and hostname in snapshot
RESOLV_TMP="$VM_DIR/resolv.conf"
HOSTNAME_TMP="$VM_DIR/hostname"
HOSTS_TMP="$VM_DIR/hosts"
printf 'nameserver %s\n' "$DNS_SERVER" >"$RESOLV_TMP"
printf '%s\n' "$(jq -r '.name // empty' "$META_JSON")" >"$HOSTNAME_TMP"
printf '127.0.0.1 localhost\n127.0.1.1 %s\n' "$(jq -r '.name // empty' "$META_JSON")" >"$HOSTS_TMP"
sudo e2rm "$DM_DEV:/etc/resolv.conf" >/dev/null 2>&1 || true
sudo e2rm "$DM_DEV:/etc/hostname" >/dev/null 2>&1 || true
sudo e2rm "$DM_DEV:/etc/hosts" >/dev/null 2>&1 || true
sudo e2cp "$RESOLV_TMP" "$DM_DEV:/etc/resolv.conf" >/dev/null 2>&1 || true
sudo e2cp "$HOSTNAME_TMP" "$DM_DEV:/etc/hostname" >/dev/null 2>&1 || true
sudo e2cp "$HOSTS_TMP" "$DM_DEV:/etc/hosts" >/dev/null 2>&1 || true

# 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

log "starting firecracker process"
rm -f "$API_SOCK"
nohup sudo -E "$FC_BIN" --api-sock "$API_SOCK" >"$LOG_FILE" 2>&1 &
FC_PID="$!"

log "waiting for firecracker api socket"
for _ in $(seq 1 200); do
  [[ -S "$API_SOCK" ]] && break
  sleep 0.02
done
[[ -S "$API_SOCK" ]] || { log "firecracker api socket not ready"; exit 1; }

log "configuring machine"
/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/machine-config \
  -H "Content-Type: application/json" \
  -d "$(jq -c '."machine-config"' "$VM_JSON")" >/dev/null

boot_args="$(jq -r '."boot-source".boot_args' "$VM_JSON")"
boot_args="$(printf '%s' "$boot_args" | sed -E 's/(^| )hostname=[^ ]+//g; s/(^| )ip=[^ ]+//g' | awk '{$1=$1; print}')"
boot_args="$boot_args ip=${GUEST_IP}::${BR_IP}:255.255.255.0::eth0:off:${DNS_SERVER}"
boot_args="$boot_args hostname=$(jq -r '.name // empty' "$META_JSON")"

log "configuring boot source"
/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/boot-source \
  -H "Content-Type: application/json" \
  -d "{
    \"kernel_image_path\": \"$KERNEL\",
    \"boot_args\": \"$boot_args\"
  }" >/dev/null

log "attaching drives"
/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/drives/rootfs \
  -H "Content-Type: application/json" \
  -d "{
    \"drive_id\": \"rootfs\",
    \"path_on_host\": \"$DM_DEV\",
    \"is_root_device\": true,
    \"is_read_only\": false
  }" >/dev/null

/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/drives/home \
  -H "Content-Type: application/json" \
  -d "{
    \"drive_id\": \"home\",
    \"path_on_host\": \"$HOME_PATH\",
    \"is_root_device\": false,
    \"is_read_only\": false
  }" >/dev/null

/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/drives/var \
  -H "Content-Type: application/json" \
  -d "{
    \"drive_id\": \"var\",
    \"path_on_host\": \"$VAR_PATH\",
    \"is_root_device\": false,
    \"is_read_only\": false
  }" >/dev/null

log "configuring network interface"
/usr/bin/sudo /usr/bin/curl --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

log "starting virtual machine"
/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -X PUT http://localhost/actions \
  -H "Content-Type: application/json" \
  -d '{ "action_type": "InstanceStart" }' >/dev/null
VM_STARTED=1

CREATED_AT="$(date -Iseconds)"
jq \
  --arg pid "$FC_PID" \
  --arg created_at "$CREATED_AT" \
  '.pid=$pid | .created_at=$created_at' \
  "$META_JSON" > "$META_JSON.tmp" && mv "$META_JSON.tmp" "$META_JSON"

VM_CONFIG_JSON="$(/usr/bin/sudo /usr/bin/curl --unix-socket "$API_SOCK" -sS http://localhost/vm/config)"
jq '.' <<<"$VM_CONFIG_JSON" > "$VM_JSON"

log "restored"