#!/bin/bash
# ==============================================================================================
# ============================= SSH Setup ======================================================
# ==============================================================================================
#
# PURPOSE
# ─────────────────────────────────────────────────────────────────────────────
# Generates the SSH keypair for rsync automation and installs it on the remote server.
# Every cross-host operation in the ecosystem — rsync, conf sync, fallback container
# control, the upgrade webhook — authenticates with this key. If it is missing or broken,
# the mesh silently degrades to single-host.
#
# Key named after this server: hostname lowercased, unraid- prefix stripped.
#   unRAID-Gmer4Lfe  → gmer4lfe_rsync_automation
#   unRAID-Jayred365 → jayred365_rsync_automation
#
# Idempotent — skips generation if the key already exists (--force to regenerate).
# Updates host*.conf with the key path on success.
#
# ==============================================================================================
# OPERATIONAL MODEL
# ==============================================================================================
#
# Setup (default):
#   1. Key exists? → skip generation unless --force
#   2. Generate keypair, named from this host
#   3. Copy the public key to the remote's authorized_keys
#   4. Verify authentication actually works before claiming success
#   5. Write the key path into host*.conf
#
# Validate (--validate), called during partnership --check cycles:
#   Remote unreachable on Tailscale → network problem, NOT counted as a strike
#   Remote reachable but SSH auth fails → key problem, strike incremented
#   Clean connectivity for SSH_STRIKE_RESET_HRS → strikes reset automatically
#   At SSH_MAX_STRIKES → notify and return exit 2 so the caller can escalate
#
#   State: DATA_DIR/ssh_strikes_{REMOTE_SERVER_NAME}.db
#
# ==============================================================================================
# DESIGN PRINCIPLES
# ==============================================================================================
#
# Distinguish Unreachable From Unauthorised
#   The strike system counts SSH auth failures only. A partner that is simply offline is a
#   network condition, and counting it would fire a key-rotation alarm every time the remote
#   reboots. Only "I can reach you but you will not let me in" is a key problem.
#
# Idempotent by Default, Destructive Only on Request
#   A bare run never replaces an existing key. Regenerating invalidates every authorized_keys
#   entry the old key was in — including on hosts this script is not talking to right now —
#   so it requires --force explicitly.
#
# Verify Before Recording
#   The key path is written into host*.conf only after authentication has been proven to
#   work. Recording a key that does not authenticate would leave every downstream script
#   pointed at a credential that silently fails.
#
# Strikes Reset on Recovery
#   Counters clear themselves after a period of clean connectivity, so a transient outage
#   does not accumulate toward an alarm across unrelated weeks.
#
# ==============================================================================================
# OPERATIONAL SAFEGUARDS
# ==============================================================================================
#
# Root Enforcement
#   Reads and writes /root/.ssh and installs keys on the remote as root.
#
# Lock Acquisition
#   acquire_lock prevents concurrent runs. Two instances generating or copying keys at once
#   could leave authorized_keys holding a key whose private half was already replaced.
#
# Host Detection
#   detect_hosts() resolves MY_ID / REMOTE_ID for key naming and remote targeting.
#
# Existing Key Protection
#   Generation is skipped when a key is present. Overwriting requires --force.
#
# Auth Verified Before Conf Write
#   host*.conf is updated only after a successful authentication test.
#
# Network-vs-Auth Discrimination
#   Unreachable remotes never increment the strike counter — see Design Principles.
#
# Strike Ceiling
#   SSH_MAX_STRIKES bounds how long a genuinely broken key goes unreported, and exit 2 lets
#   the caller decide whether that is escalation-worthy.
#
# Local-Only Escape Hatch
#   --local-only generates the key without touching the remote, for onboarding a partner
#   that is not reachable yet.
#
# Dry Run Support
#   --dry-run previews generation, copy and conf update without performing any.
#
# ==============================================================================================
# CONFIGURATION
# ==============================================================================================
#
# master.conf
#
#   SSH_MAX_STRIKES
#       Consecutive SSH auth failures before notifying (default: 5)
#
#   SSH_STRIKE_RESET_HRS
#       Hours of clean connectivity before the strike counter resets (default: 24)
#
# host*.conf
#
#   HOST*_SSH_KEY
#       Written by this script on success. Read by rsync.sh, conf_sync.sh, fallback.sh and
#       upgrade_webhook_handler.sh — every cross-host operation depends on it.
#
# ==============================================================================================
# RUNTIME MODES
# ==============================================================================================
#
# ssh_setup.sh
#     Initial setup — generate if missing, copy to remote, update conf. Idempotent.
#
# ssh_setup.sh --force
#     Regenerate the key even if one exists, and re-copy to the remote.
#
# ssh_setup.sh --validate
#     Health check with strike tracking. Exit 2 at the strike limit.
#
# ssh_setup.sh --status
#     Show key state, fingerprint, and remote connectivity. Then exit.
#
# ssh_setup.sh --local-only
#     Generate the key locally and skip the remote copy.
#
# ssh_setup.sh --dry-run / --log
#     Supported by every mode above.
#
# ==============================================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPTS_ROOT="$SCRIPT_DIR/.."

source "$SCRIPTS_ROOT/load_config.sh"

# ── Parse --force before parse_args ───────────────────────────────────────────────────────────
MODE="setup"
FORCE=false
LOCAL_ONLY=false
FILTERED_ARGS=()

for arg in "$@"; do
    case "$arg" in
        --force)      FORCE=true ;;
        --validate)   MODE="validate" ;;
        --status)     MODE="status" ;;
        --local-only) LOCAL_ONLY=true ;;
        *)            FILTERED_ARGS+=("$arg") ;;
    esac
done

parse_args "${FILTERED_ARGS[@]}"

# ==============================================================================================
# ━━━ Setup ━━━
# ==============================================================================================
if [[ "$EUID" -ne 0 ]]; then
    error "Must be run as root"
    exit 1
fi

acquire_lock

detect_hosts

# ── Derive key name from hostname ─────────────────────────────────────────────────────────────
# derive_short_name() title-cases the result (for display elsewhere) — lowercase it here,
# same as before, since this feeds a filename → gmer4lfe_rsync_automation
SHORT_NAME="$(derive_short_name "$LOCAL_SERVER_NAME")"
SHORT_NAME="${SHORT_NAME,,}"
SSH_KEY_NAME="${SHORT_NAME}_rsync_automation"
SSH_KEY_PATH="/root/.ssh/${SSH_KEY_NAME}"
SSH_PUB_PATH="${SSH_KEY_PATH}.pub"

# ── Host conf path ────────────────────────────────────────────────────────────────────────────
HOST_NUM="${MY_ID#HOST}"   # "1" or "2"
HOST_CONF="$SCRIPTS_ROOT/Configurations/host${HOST_NUM}.conf"
KEY_CONF_VAR="${MY_ID}_SSH_KEY"

# ── Strike state file ─────────────────────────────────────────────────────────────────────────
SSH_STRIKE_FILE="$DATA_DIR/ssh_strikes_${REMOTE_SERVER_NAME}.db"

# ==============================================================================================
# ─── HELPERS ──────────────────────────────────────────────────────────────────────────────────
# ==============================================================================================

test_ssh_auth() {
    local remote_ip="$1"
    timeout "${SSH_TIMEOUT:-15}" ssh \
        -i "$SSH_KEY_PATH" \
        -o BatchMode=yes \
        -o ConnectTimeout="${SSH_TIMEOUT:-15}" \
        -o StrictHostKeyChecking=no \
        root@"$remote_ip" "echo ok" >/dev/null 2>&1
}

read_strikes() {
    grep "^strikes=" "$SSH_STRIKE_FILE" 2>/dev/null | cut -d= -f2 || echo 0
}

read_last_strike_epoch() {
    local ts
    ts=$(grep "^last_strike=" "$SSH_STRIKE_FILE" 2>/dev/null | cut -d= -f2)
    [[ -n "$ts" ]] && date -d "$ts" +%s 2>/dev/null || echo 0
}

write_strike_file() {
    local strikes="$1" last_strike="${2:-}" last_success="${3:-}"
    mkdir -p "$(dirname "$SSH_STRIKE_FILE")"
    cat > "$SSH_STRIKE_FILE" <<EOF
strikes=${strikes}
last_strike=${last_strike}
last_success=${last_success}
updated=$(date '+%Y-%m-%d %H:%M:%S')
EOF
}

update_conf_key_path() {
    if [[ ! -f "$HOST_CONF" ]]; then
        warn "$(basename "$HOST_CONF") not found — update ${KEY_CONF_VAR} manually to: $SSH_KEY_PATH"
        return 1
    fi

    local current
    current=$(grep "^[[:space:]]*${KEY_CONF_VAR}=" "$HOST_CONF" 2>/dev/null | \
              sed 's/.*="\?\([^"]*\)"\?.*/\1/')

    if [[ "$current" == "$SSH_KEY_PATH" ]]; then
        log "${KEY_CONF_VAR} already correct in $(basename "$HOST_CONF")"
        return 0
    fi

    if grep -q "^[[:space:]]*${KEY_CONF_VAR}=" "$HOST_CONF" 2>/dev/null; then
        sed -i "s|^[[:space:]]*${KEY_CONF_VAR}=.*|    ${KEY_CONF_VAR}=\"${SSH_KEY_PATH}\"|" "$HOST_CONF" && \
            echo "${KEY_CONF_VAR} updated in $(basename "$HOST_CONF") ✅" || \
            warn "Failed to update ${KEY_CONF_VAR} in $(basename "$HOST_CONF") — update manually"
    else
        warn "${KEY_CONF_VAR} not found in $(basename "$HOST_CONF") — add manually:"
        warn "    ${KEY_CONF_VAR}=\"${SSH_KEY_PATH}\""
    fi
}

# ==============================================================================================
# ━━━ Status ━━━
# ==============================================================================================
if [[ "$MODE" == "status" ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY SSH SETUP STATUS ━━━━━"
    echo "$ICON_HOST  Identity: $MY_ID ($LOCAL_SERVER_NAME)"
    echo "$ICON_NET   Remote:   $REMOTE_SERVER_NAME"
    echo ""
    echo "  Key name: $SSH_KEY_NAME"
    echo "  Key path: $SSH_KEY_PATH"
    echo ""

    if [[ -f "$SSH_KEY_PATH" ]]; then
        local_fp=$(ssh-keygen -lf "$SSH_KEY_PATH" 2>/dev/null || echo "unreadable")
        echo "  $ICON_DONE  Private key: present"
        echo "         Fingerprint: $local_fp"
    else
        echo "  $ICON_ERROR Private key: NOT found"
    fi

    if [[ -f "$SSH_PUB_PATH" ]]; then
        echo "  $ICON_DONE  Public key:  present"
    else
        echo "  $ICON_ERROR Public key:  NOT found"
    fi

    # Check conf
    if [[ -f "$HOST_CONF" ]]; then
        current_conf=$(grep "^[[:space:]]*${KEY_CONF_VAR}=" "$HOST_CONF" 2>/dev/null | \
                       sed 's/.*="\?\([^"]*\)"\?.*/\1/')
        if [[ "$current_conf" == "$SSH_KEY_PATH" ]]; then
            echo "  $ICON_DONE  Conf ($(basename "$HOST_CONF")): ${KEY_CONF_VAR} ✅"
        else
            echo "  $ICON_ERROR Conf ($(basename "$HOST_CONF")): ${KEY_CONF_VAR}=${current_conf:-not set}"
        fi
    fi

    # Remote connectivity
    echo ""
    REMOTE_IP=$(resolve_tailscale_ip "$REMOTE_SERVER_NAME")
    if [[ -z "$REMOTE_IP" ]]; then
        echo "  $ICON_ERROR Remote ($REMOTE_SERVER_NAME): Tailscale unreachable"
    elif [[ -f "$SSH_KEY_PATH" ]] && test_ssh_auth "$REMOTE_IP"; then
        echo "  $ICON_DONE  Remote ($REMOTE_SERVER_NAME [$REMOTE_IP]): SSH auth OK ✅"
    else
        echo "  $ICON_ERROR Remote ($REMOTE_SERVER_NAME [$REMOTE_IP]): SSH auth FAILED"
    fi

    # Strike state
    if [[ -f "$SSH_STRIKE_FILE" ]]; then
        echo ""
        echo "  Strike state:"
        while IFS= read -r line; do
            echo "    $line"
        done < "$SSH_STRIKE_FILE"
    fi

    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

# ==============================================================================================
# ━━━ Validate (SSH health check + strike tracking) ━━━
# ==============================================================================================
if [[ "$MODE" == "validate" ]]; then
    MAX_STRIKES="${SSH_MAX_STRIKES:-5}"
    RESET_HRS="${SSH_STRIKE_RESET_HRS:-24}"

    REMOTE_IP=$(resolve_tailscale_ip "$REMOTE_SERVER_NAME")
    if [[ -z "$REMOTE_IP" ]]; then
        log "SSH validate — $REMOTE_SERVER_NAME Tailscale unreachable, not an SSH issue"
        exit 0
    fi

    if [[ ! -f "$SSH_KEY_PATH" ]]; then
        warn "SSH validate — key not found at $SSH_KEY_PATH — run ssh_setup.sh to create"
        exit 1
    fi

    NOW=$(date '+%Y-%m-%d %H:%M:%S')
    NOW_EPOCH=$(date +%s)

    if test_ssh_auth "$REMOTE_IP"; then
        # SSH works — reset strikes if they were > 0
        STRIKES=$(read_strikes)
        if [[ "$STRIKES" -gt 0 ]]; then
            write_strike_file 0 "" "$NOW"
            echo "SSH validate — auth restored to $REMOTE_SERVER_NAME ✅ (strikes reset)"
        else
            echo "SSH validate — $REMOTE_SERVER_NAME SSH auth OK ✅"
        fi
        exit 0
    fi

    # SSH auth failed — check if strikes should reset first
    STRIKES=$(read_strikes)
    LAST_STRIKE_EPOCH=$(read_last_strike_epoch)
    RESET_SECS=$(( RESET_HRS * 3600 ))

    if [[ "$STRIKES" -gt 0 ]] && \
       [[ "$LAST_STRIKE_EPOCH" -gt 0 ]] && \
       [[ $(( NOW_EPOCH - LAST_STRIKE_EPOCH )) -gt "$RESET_SECS" ]]; then
        warn "SSH validate — strike counter reset after ${RESET_HRS}hr gap"
        STRIKES=0
    fi

    STRIKES=$(( STRIKES + 1 ))
    write_strike_file "$STRIKES" "$NOW" ""
    warn "SSH validate — $REMOTE_SERVER_NAME auth FAILED (strike $STRIKES/$MAX_STRIKES)"

    if [[ "$STRIKES" -ge "$MAX_STRIKES" ]]; then
        error "SSH auth to $REMOTE_SERVER_NAME broken — ${STRIKES} consecutive failures"
        error "Repair: run 'Partnership/ssh_setup.sh --force' to regenerate and re-copy key"
        notify "SSH auth broken to $REMOTE_SERVER_NAME ($MY_ID) — ${STRIKES} strikes, manual repair needed. Run: Partnership/ssh_setup.sh --force" \
            "SSH Setup" "warning"
        exit 2
    fi

    exit 1
fi

# ==============================================================================================
# ━━━ Setup (default) or --force ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_GEAR SSH Key Setup — $MY_ID ($LOCAL_SERVER_NAME) — $(date '+%Y-%m-%d %H:%M:%S') ━━━"
echo ""
echo "  Key name: $SSH_KEY_NAME"
echo "  Key path: $SSH_KEY_PATH"

[[ "$DRY_RUN" == true ]] && warn "DRY RUN — no keys will be created or copied"

# ── Key generation ────────────────────────────────────────────────────────────────────────────
echo ""
echo "━━━ $ICON_GEAR Key Generation ━━━"

if [[ -f "$SSH_KEY_PATH" ]] && [[ "$FORCE" == false ]]; then
    local_fp=$(ssh-keygen -lf "$SSH_KEY_PATH" 2>/dev/null || echo "unreadable")
    echo "Key already exists — skipping generation (--force to regenerate)"
    echo "  $local_fp"
else
    if [[ "$FORCE" == true ]] && [[ -f "$SSH_KEY_PATH" ]]; then
        warn "Regenerating key (--force) — existing key will be replaced"
        if [[ "$DRY_RUN" == false ]]; then
            rm -f "$SSH_KEY_PATH" "$SSH_PUB_PATH"
        fi
    fi

    warn "Generating ed25519 keypair: $SSH_KEY_PATH"
    if [[ "$DRY_RUN" == false ]]; then
        ssh-keygen -t ed25519 -N "" -f "$SSH_KEY_PATH" -C "${SSH_KEY_NAME}@${LOCAL_SERVER_NAME}" && \
            echo "Keypair generated ✅" || {
            error "Failed to generate keypair"
            exit 1
        }
        chmod 600 "$SSH_KEY_PATH"
        chmod 644 "$SSH_PUB_PATH"
    else
        warn "DRY RUN — would generate: $SSH_KEY_PATH"
    fi
fi

# ── Update conf ───────────────────────────────────────────────────────────────────────────────
echo ""
echo "━━━ $ICON_GEAR Update host${HOST_NUM}.conf ━━━"

if [[ "$DRY_RUN" == false ]]; then
    update_conf_key_path
else
    warn "DRY RUN — would set ${KEY_CONF_VAR}=\"${SSH_KEY_PATH}\" in host${HOST_NUM}.conf"
fi

# ── Copy to remote ────────────────────────────────────────────────────────────────────────────
if [[ "$LOCAL_ONLY" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY SSH SETUP SUMMARY (local) ━━━━━"
    echo "  Key:    $SSH_KEY_PATH"
    echo "  Pubkey: $SSH_PUB_PATH"
    echo "  Conf:   ${KEY_CONF_VAR} in host${HOST_NUM}.conf"
    echo ""
    warn "Local setup complete — copy public key to remote manually:"
    warn "  ssh-copy-id -i $SSH_PUB_PATH root@<remote>"
    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

echo ""
echo "━━━ $ICON_NET Copy Public Key to Remote ($REMOTE_SERVER_NAME) ━━━"

resolve_remote_ip

if [[ "$DRY_RUN" == true ]]; then
    warn "DRY RUN — would ssh-copy-id to root@$REMOTE_SERVER"
    warn "DRY RUN — would require password for root@$REMOTE_SERVER if key auth not yet set up"
else
    warn "Installing public key on $REMOTE_SERVER_NAME ($REMOTE_SERVER)..."
    warn "(Password prompt for root@$REMOTE_SERVER is expected on first setup)"

    if ssh-copy-id -i "$SSH_PUB_PATH" -o ConnectTimeout="${SSH_TIMEOUT:-15}" \
        root@"$REMOTE_SERVER" 2>/dev/null; then
        echo "Public key installed on $REMOTE_SERVER_NAME ✅"
    else
        error "ssh-copy-id failed — check that:"
        error "  1. Remote server is reachable: tailscale status"
        error "  2. Password auth is enabled on remote: grep PasswordAuthentication /etc/ssh/sshd_config"
        error "  3. The correct password for root@$REMOTE_SERVER_NAME is used"
        exit 1
    fi
fi

# ── Verify ────────────────────────────────────────────────────────────────────────────────────
echo ""
echo "━━━ $ICON_VERIFY Verify SSH Auth ━━━"

if [[ "$DRY_RUN" == false ]]; then
    if test_ssh_auth "$REMOTE_SERVER"; then
        echo "SSH auth to $REMOTE_SERVER_NAME working ✅"
        if [[ -f "$SSH_STRIKE_FILE" ]]; then
            write_strike_file 0 "" "$(date '+%Y-%m-%d %H:%M:%S')"
        fi
    else
        warn "SSH auth test failed — key may need a moment to propagate"
        warn "Verify manually: ssh -i $SSH_KEY_PATH root@$REMOTE_SERVER"
    fi
else
    warn "DRY RUN — would test: ssh -i $SSH_KEY_PATH root@$REMOTE_SERVER echo ok"
fi

# ── Summary ───────────────────────────────────────────────────────────────────────────────────
echo ""
echo "━━━━━ $ICON_SUMMARY SSH SETUP SUMMARY ━━━━━"
echo "  Key:    $SSH_KEY_PATH"
echo "  Remote: $REMOTE_SERVER_NAME ($REMOTE_SERVER)"
echo "  Conf:   ${KEY_CONF_VAR} in host${HOST_NUM}.conf"
echo ""
[[ "$DRY_RUN" == true ]] && warn "DRY RUN — no changes made" || \
    warn "$ICON_DONE DONE — SSH key ready for rsync automation ✅"
echo "━━━━━━━━━━━━━━━━━━━━━━━"
exit 0
