#!/bin/bash
# ==============================================================================================
# ============================= Onboard Cancel =================================================
# ==============================================================================================
#
# PURPOSE
# ─────────────────────────────────────────────────────────────────────────────
# Removes SSH keys between HOST1 and HOST2 in the specified direction.
# Safe to run at any phase. Clears related setup.db flags.
#
# The escape hatch for a half-finished onboard: partnership setup is multi-phase, and an
# attempt abandoned midway leaves keys installed and phase flags set. This unwinds that so
# onboarding can be started cleanly rather than resumed from an unknown state.
#
# ==============================================================================================
# OPERATIONAL MODEL
# ==============================================================================================
#
# Per direction requested:
#
#   h1 (HOST1 → HOST2)
#     1. SSH to HOST2 and delete HOST1's public key line from its authorized_keys
#     2. Clear HOST2's phase / key-ready flags from the setup db
#     3. Delete the local HOST1 key pair
#
#   h2 (HOST2 → HOST1)
#     Remove HOST2's public key from HOST1's authorized_keys, identifying the key by
#     HOST2's hostname in the key comment.
#
#   both — run each direction in turn.
#
# Key removal is matched on the key blob or hostname comment, never on line number.
#
# ==============================================================================================
# DESIGN PRINCIPLES
# ==============================================================================================
#
# Safe at Any Phase
#   Onboarding is multi-phase and can fail anywhere in it. This runs against whatever state
#   exists rather than requiring a known starting point — a missing key or an absent flag is
#   a no-op, not an error. Cancelling twice is harmless.
#
# Direction Is Explicit
#   Removing a key is not symmetric: it breaks authentication for whichever side loses it.
#   The direction must be stated, and defaults to h1 (this host's own outbound key) rather
#   than to both, so an unqualified run cannot sever the partner's access to you.
#
# Unwind, Do Not Repair
#   The job is to return to a clean pre-onboard state so onboarding can be re-run from the
#   top. It deliberately does not attempt to salvage or resume a partial setup — a known
#   empty state is worth more than a guessed-at partial one.
#
# Keys and Flags Together
#   Removing the key without clearing the setup-db flags would leave onboarding believing a
#   phase had completed. Both are cleared in the same pass for that reason.
#
# ==============================================================================================
# OPERATIONAL SAFEGUARDS
# ==============================================================================================
#
# Root Enforcement
#   Deletes key pairs from /root/.ssh and edits authorized_keys on both ends as root.
#
# Lock Acquisition
#   acquire_lock prevents this racing an in-progress onboard, which would otherwise be
#   installing the very keys this is removing.
#
# Host Detection
#   detect_hosts() resolves MY_ID / REMOTE_ID so the direction flags map to real hosts.
#
# Direction Default
#   Defaults to h1 — never removes the partner's inbound key unless explicitly asked.
#
# Targeted Key Removal
#   authorized_keys lines are matched by key blob or hostname comment. Nothing is removed
#   positionally, so an unrelated key can never be deleted because it sat on a given line.
#
# SSH Timeout
#   The remote edit is wrapped in SSH_TIMEOUT — an unreachable partner fails fast rather
#   than hanging a cancel that still has local cleanup to do.
#
# Idempotent
#   Absent keys and absent flags are skipped silently. Re-running is safe.
#
# Dry Run Support
#   --dry-run reports every key and flag it would remove, and removes none.
#
# ==============================================================================================
# CONFIGURATION
# ==============================================================================================
#
# master.conf
#
#   SSH_TIMEOUT
#       Bounds the remote authorized_keys edit.
#
# host*.conf
#
#   SSH_KEY
#       Key used to reach the partner, and the local pair deleted in the h1 direction.
#
#   HOST* — hostnames, used to identify which key comment belongs to which side.
#
# Setup state lives in the platform setup db (platform_setup_db_path) — the phase and
# key-ready flags cleared here are the same ones partnership_onboard.sh sets.
#
# ==============================================================================================
# 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="$(platform_setup_db_path)"
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\" $(platform_setup_db_path) 2>/dev/null
                 echo ok" 2>/dev/null | grep -q ok && {
                echo "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" && echo "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"
            echo "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" && {
            echo "$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
