#!/bin/bash
# Restore OpenClaw VM from backup
# Usage: restore-openclaw-vm.sh [instance] [target-host]
#   instance    - registry name (default: zap)
#   target-host - override host IP for fresh VM deployments

set -euo pipefail

INSTANCE="${1:-zap}"
TARGET_HOST="${2:-}"
REGISTRY="${HOME}/.claude/state/openclaw-instances.json"
ANSIBLE_DIR="${HOME}/lab/swarm/ansible"

# Resolve instance from registry
GUEST_USER=$(python3 -c "
import json
data = json.load(open('${REGISTRY}'))
inst = next((i for i in data['instances'] if i['name'] == '${INSTANCE}'), None)
if not inst: print('openclaw')
else: print(inst['user'])
")

if [[ -z "${TARGET_HOST}" ]]; then
    TARGET_HOST=$(python3 -c "
import json, sys
data = json.load(open('${REGISTRY}'))
inst = next((i for i in data['instances'] if i['name'] == '${INSTANCE}'), None)
if not inst or not inst.get('host'):
    print('ERROR: no host in registry — pass target-host as second argument', file=sys.stderr)
    sys.exit(1)
print(inst['host'])
")
fi

SYNC_DIR="${HOME}/lab/swarm/openclaw"
MINIO_BUCKET="zap"
MINIO_PREFIX="backups"

# Find config source — local sync dir first, fall back to MinIO archive
USE_MINIO=false
if [[ ! -d "${SYNC_DIR}" || -z "$(ls -A "${SYNC_DIR}" 2>/dev/null)" ]]; then
    echo "No local sync found in ${SYNC_DIR}, fetching latest archive from MinIO..."
    LATEST_KEY=$(aws s3 ls "s3://${MINIO_BUCKET}/${MINIO_PREFIX}/${INSTANCE}-" \
        | awk '{print $4}' | sort | tail -1)
    if [[ -z "${LATEST_KEY}" ]]; then
        echo "ERROR: No config found locally or in s3://${MINIO_BUCKET}/${MINIO_PREFIX}/" >&2
        exit 1
    fi
    MINIO_ARCHIVE=$(mktemp --suffix=".tar.gz")
    trap 'rm -f "${MINIO_ARCHIVE}"' EXIT
    aws s3 cp "s3://${MINIO_BUCKET}/${MINIO_PREFIX}/${LATEST_KEY}" "${MINIO_ARCHIVE}"
    echo "Downloaded: ${LATEST_KEY}"
    USE_MINIO=true
fi

SOURCE_DESC="${SYNC_DIR}"
[[ "${USE_MINIO}" == true ]] && SOURCE_DESC="MinIO: ${LATEST_KEY}"

echo "========================================"
echo " OpenClaw VM Restore"
echo "========================================"
echo " Instance : ${INSTANCE}"
echo " Target   : ${TARGET_HOST}"
echo " Source   : ${SOURCE_DESC}"
echo "========================================"
echo ""
read -rp "Proceed? [y/N] " confirm
[[ "${confirm}" =~ ^[Yy]$ ]] || { echo "Aborted."; exit 0; }

cd "${ANSIBLE_DIR}"

# Step 1: Provision VM if a target host override was given (implies fresh VM needed)
if [[ -n "${2:-}" ]]; then
    echo ""
    echo "==> [1/4] Provisioning VM (libvirt)..."
    ansible-playbook -i inventory.yml playbooks/provision-vm.yml --limit "${INSTANCE}" \
        -e "vm_ip=${TARGET_HOST}" -e "ansible_host=${TARGET_HOST}"
else
    echo ""
    echo "==> [1/4] Skipping VM provisioning (using existing VM at ${TARGET_HOST})"
fi

echo ""
echo "==> [2/4] Provisioning guest OS..."
ansible-playbook -i inventory.yml playbooks/install.yml --limit "${INSTANCE}" \
    -e "ansible_host=${TARGET_HOST}"

echo ""
echo "==> [3/4] Applying customizations..."
ansible-playbook -i inventory.yml playbooks/customize.yml --limit "${INSTANCE}" \
    -e "ansible_host=${TARGET_HOST}"

# Step 3: Restore config
echo ""
echo "==> [4/5] Restoring OpenClaw config..."

# Stop service before restoring
ssh -o StrictHostKeyChecking=no "root@${TARGET_HOST}" \
    "su - ${GUEST_USER} -c 'systemctl --user stop openclaw-gateway.service 2>/dev/null || true'"

# Push config and restore
if [[ "${USE_MINIO}" == true ]]; then
    cat "${MINIO_ARCHIVE}" | ssh -o StrictHostKeyChecking=no "root@${TARGET_HOST}" \
        "cd /home/${GUEST_USER} && tar xzf - && chown -R ${GUEST_USER}:${GUEST_USER} openclaw"
else
    rsync -az \
        -e "ssh -o StrictHostKeyChecking=no -o BatchMode=yes" \
        "${SYNC_DIR}/" \
        "root@${TARGET_HOST}:/home/${GUEST_USER}/.openclaw/"
    ssh -o StrictHostKeyChecking=no "root@${TARGET_HOST}" \
        "chown -R ${GUEST_USER}:${GUEST_USER} /home/${GUEST_USER}/.openclaw"
fi

echo "    Config restored from: ${SOURCE_DESC}"

# Step 4: Start service
echo ""
echo "==> [5/5] Starting service..."
ssh -o StrictHostKeyChecking=no "root@${TARGET_HOST}" \
    "su - ${GUEST_USER} -c 'systemctl --user daemon-reload && systemctl --user start openclaw-gateway.service'"

# Verify
sleep 3
STATUS=$(ssh -o StrictHostKeyChecking=no "root@${TARGET_HOST}" \
    "su - ${GUEST_USER} -c 'systemctl --user is-active openclaw-gateway.service 2>/dev/null'" || echo "unknown")

echo ""
echo "========================================"
echo " Restore complete"
echo " Service status: ${STATUS}"
echo "========================================"