#!/bin/bash
# ==============================================================================================
# ============================= SSH Setup ======================================================
# ==============================================================================================
# Generates the SSH keypair for rsync automation and installs it on the remote server.
# Key named after this server: hostname lowercased, unraid- prefix stripped.
#   unRAID-Gmer4Lfe → gmer4lfe_rsync_automation
#   unRAID-Jayred365 → jayred365_rsync_automation
# Idempotent — skips generation if key already exists (use --force to regenerate).
# Updates host*.conf with key path on success.
#
# ── MODES ─────────────────────────────────────────────────────────────────────────────────────
#   (default)       — generate key if missing, copy to remote, update conf
#   --force         — regenerate key even if it exists, re-copy to remote
#   --validate      — test SSH auth to remote, track strikes, notify at limit
#   --status        — show key state, fingerprint, remote connectivity
#   --dry-run       — preview without creating, copying, or updating conf
#
# ── STRIKE SYSTEM (--validate) ────────────────────────────────────────────────────────────────
# Called during partnership --check cycles to detect broken SSH auth.
# Tracks consecutive SSH auth failures — not network unreachability.
# Remote Tailscale IP unreachable = network issue → not counted as SSH strike.
# Remote reachable but SSH auth fails = key issue → strike incremented.
# Strikes reset automatically after SSH_STRIKE_RESET_HRS of clean connectivity.
# At SSH_MAX_STRIKES: notify + return exit 2 (caller can escalate).
# State: DATA_DIR/ssh_strikes_{REMOTE_SERVER_NAME}.db
#
# ── CONFIGURATION (master.conf) ───────────────────────────────────────────────────────────────
#   SSH_MAX_STRIKES       — consecutive failures before notifying (default 5)
#   SSH_STRIKE_RESET_HRS  — hours since last failure before counter resets (default 24)
#
# ── USAGE ─────────────────────────────────────────────────────────────────────────────────────
#   Partnership/ssh_setup.sh                — initial setup (idempotent)
#   Partnership/ssh_setup.sh --force        — regenerate + re-copy
#   Partnership/ssh_setup.sh --validate     — health check + strike tracking
#   Partnership/ssh_setup.sh --status       — show key and connectivity state
#   Partnership/ssh_setup.sh --local-only   — generate key locally, skip remote copy
#   Any mode supports --dry-run and --log
# ==============================================================================================

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 ─────────────────────────────────────────────────────────────
# Strip unraid- prefix (case-insensitive) if present → gmer4lfe_rsync_automation
SHORT_NAME="${LOCAL_SERVER_NAME,,}"
[[ "${SHORT_NAME}" == unraid-* ]] && SHORT_NAME="${SHORT_NAME:7}"
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" && \
            log "${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"
            log "SSH validate — auth restored to $REMOTE_SERVER_NAME ✅ (strikes reset)"
        else
            log "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}" && \
            log "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