#!/bin/bash
# ==============================================================================================
# ============================= Delete Keys ====================================================
# ==============================================================================================
#
# PURPOSE
# ─────────────────────────────────────────────────────────────────────────────
# Removes SSH keys between HOST1 and HOST2 in the specified direction.
# Safe to run at any phase. Clears related setup.db flags.
#
# ==============================================================================================
# RUNTIME MODES
# ==============================================================================================
#
# Partnership/onboard_cancel.sh --direction=h1   (default)
#     HOST1 → HOST2: remove HOST1's public key from HOST2's authorized_keys,
#     delete local HOST1 key pair, clear HOST2 phase flags from setup.db.
#
# Partnership/onboard_cancel.sh --direction=h2
#     HOST2 → HOST1: remove HOST2's public key from HOST1's authorized_keys.
#     Identifies the key by HOST2's hostname in the key comment.
#
# Partnership/onboard_cancel.sh --direction=both
#     Both directions.
#
# Partnership/onboard_cancel.sh --dry-run
#     Preview without making changes.
#
# ==============================================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPTS_ROOT="$SCRIPT_DIR/.."
SSH_TIMEOUT=15
DIRECTION="h1"
FILTERED_ARGS=()

for arg in "$@"; do
    case "$arg" in
        --direction=*) DIRECTION="${arg#--direction=}" ;;
        *)             FILTERED_ARGS+=("$arg") ;;
    esac
done

source "$SCRIPTS_ROOT/load_config.sh"
parse_args "${FILTERED_ARGS[@]}"

[[ "$EUID" -ne 0 ]] && { error "Must be run as root"; exit 1; }

acquire_lock

detect_hosts

OWNER_ID="${PARTNERSHIP_OWNER_HOST:-HOST1}"
MIRROR_ID=$( [[ "$OWNER_ID" == "HOST1" ]] && echo "HOST2" || echo "HOST1" )
MIRROR="${!MIRROR_ID}"
SSH_KEY_PUB="${SSH_KEY}.pub"
STATE_FILE="${VARAVERK_SETUP_FILE:-${STATE_DIR:-/boot/config}/varaverk_setup.db}"
AUTH_KEYS="/root/.ssh/authorized_keys"
MIRROR_SHORT="${MIRROR%%.*}"

START=$(date +%s)

echo ""
echo "━━━ Delete Keys — direction:${DIRECTION} — $(date '+%Y-%m-%d %H:%M:%S') ━━━"
[[ "$DRY_RUN" == true ]] && warn "DRY RUN — no changes will be made"
echo ""

H1_DONE=false
H2_DONE=false

# ── HOST1 → HOST2: remove HOST1's key from HOST2 + delete local pair ──────────
if [[ "$DIRECTION" == "h1" || "$DIRECTION" == "both" ]]; then
    echo "━━━ HOST1 → HOST2: Remove HOST1 key from $MIRROR ━━━"

    if [[ ! -f "$SSH_KEY_PUB" ]]; then
        log "No local public key — nothing to remove from $MIRROR"
        H1_DONE=true
    else
        KEY_BLOB=$(awk '{print $2}' "$SSH_KEY_PUB")
        MIRROR_IP=$(resolve_tailscale_ip "$MIRROR" 2>/dev/null || true)

        if [[ -z "$MIRROR_IP" ]]; then
            warn "Cannot resolve $MIRROR Tailscale IP — remove HOST1 key from $MIRROR manually"
        elif [[ "$DRY_RUN" == true ]]; then
            warn "DRY RUN — would remove HOST1 key from $MIRROR:/root/.ssh/authorized_keys"
            H1_DONE=true
        else
            timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
                -o ConnectTimeout="$SSH_TIMEOUT" -o BatchMode=yes root@"$MIRROR_IP" \
                "sed -i \"|${KEY_BLOB}|d\" /root/.ssh/authorized_keys 2>/dev/null
                 sed -i \"/^${MIRROR_ID}_PHASE\|^${MIRROR_ID}_KEY_READY/d\" /boot/config/varaverk_setup.db 2>/dev/null
                 echo ok" 2>/dev/null | grep -q ok && {
                log "HOST1 key removed from $MIRROR authorized_keys ✅"
                H1_DONE=true
            } || warn "Could not SSH to $MIRROR — remove HOST1 key there manually"
        fi
    fi

    # Delete local key pair
    if [[ ! -f "$SSH_KEY" && ! -f "$SSH_KEY_PUB" ]]; then
        log "Local key already gone"
    elif [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would delete: $SSH_KEY and ${SSH_KEY}.pub"
    else
        rm -f "$SSH_KEY" "$SSH_KEY_PUB" && log "Local key pair deleted ✅" || \
            warn "Failed to delete local key — check permissions"
    fi

    # Clear HOST2 phase flags from local setup.db
    if [[ -f "$STATE_FILE" ]]; then
        if [[ "$DRY_RUN" == false ]]; then
            sed -i "/^${MIRROR_ID}_PHASE/d; /^${MIRROR_ID}_KEY_READY/d" "$STATE_FILE"
            log "Phase flags cleared from local setup.db ✅"
        else
            warn "DRY RUN — would clear ${MIRROR_ID}_PHASE* from setup.db"
        fi
    fi
    echo ""
fi

# ── HOST2 → HOST1: remove HOST2's key from HOST1's authorized_keys ────────────
if [[ "$DIRECTION" == "h2" || "$DIRECTION" == "both" ]]; then
    echo "━━━ HOST2 → HOST1: Remove $MIRROR key from HOST1 ━━━"

    if [[ ! -f "$AUTH_KEYS" ]]; then
        log "No authorized_keys on HOST1 — nothing to remove"
        H2_DONE=true
    elif ! grep -qi "$MIRROR_SHORT" "$AUTH_KEYS" 2>/dev/null; then
        log "$MIRROR key not found in HOST1 authorized_keys (already removed or never added)"
        H2_DONE=true
    elif [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would remove $MIRROR_SHORT key from $AUTH_KEYS"
        H2_DONE=true
    else
        sed -i "/${MIRROR_SHORT}/Id" "$AUTH_KEYS" && {
            log "$MIRROR key removed from HOST1 authorized_keys ✅"
            H2_DONE=true
        } || warn "Failed to remove $MIRROR key from HOST1 authorized_keys"
    fi
    echo ""
fi

# ── Summary ───────────────────────────────────────────────────────────────────
END=$(date +%s)
echo "━━━━━ DONE ━━━━━"
[[ "$DIRECTION" == "h1" || "$DIRECTION" == "both" ]] && \
    echo "  HOST1 → HOST2: $( [[ "$H1_DONE" == true ]] && echo "✅" || echo "⚠ manual step may be needed" )"
[[ "$DIRECTION" == "h2" || "$DIRECTION" == "both" ]] && \
    echo "  HOST2 → HOST1: $( [[ "$H2_DONE" == true ]] && echo "✅" || echo "⚠" )"
echo "  Duration: $(format_duration $(( END - START )))"
echo "━━━━━━━━━━━━━━━━━━━━━━━"
exit 0