#!/bin/bash
# reboot-with-check.sh — Reboot a server and wait for it + key services to come back.
# Usage: reboot-with-check.sh <hostname|ip> [service_port...]
#   reboot-with-check.sh 5.161.62.38 80 443 22
#   reboot-with-check.sh app1-bu.itpropartner.com 22

set -euo pipefail

HOST="$1"
shift
SERVICE_PORTS=("$@")
MAX_BOOT_WAIT=120   # seconds to wait for SSH after reboot
MAX_SERVICE_WAIT=30 # seconds to wait for each service port
PING_COUNT=3

log() { echo "[$(date -u +'%H:%M:%S')] $*"; }

if [ -z "$HOST" ]; then
    echo "Usage: $0 <hostname|ip> [port...]"
    exit 1
fi

log "=== Rebooting $HOST ==="

# Step 1: Trigger reboot via SSH
log "Triggering reboot..."
ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    root@"$HOST" "reboot" 2>/dev/null || true

# Step 2: Wait for server to go down
log "Waiting for server to go down..."
DOWN=0
for i in $(seq 1 15); do
    if ! ping -c 1 -W 2 "$HOST" >/dev/null 2>&1; then
        DOWN=$((DOWN + 1))
        if [ "$DOWN" -ge 2 ]; then
            log "Server is down (confirmed after ${i}s)"
            break
        fi
    fi
    sleep 2
done

# Step 3: Wait for SSH to come back
log "Waiting for SSH to come back..."
START=$(date +%s)
while true; do
    ELAPSED=$(( $(date +%s) - START ))
    if [ "$ELAPSED" -ge "$MAX_BOOT_WAIT" ]; then
        log "TIMEOUT: SSH not responding after ${MAX_BOOT_WAIT}s"
        exit 1
    fi
    if timeout 5 bash -c "echo > /dev/tcp/$HOST/22" 2>/dev/null; then
        log "SSH back after ${ELAPSED}s"
        break
    fi
    sleep 3
done

# Step 4: Wait for stable connectivity (ping-pong buffer)
sleep 5

# Step 5: Verify hostname (make sure it's NOT in rescue mode)
HOSTNAME=$(ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    -i "$HOME/.ssh/itpp-infra" root@"$HOST" "hostname" 2>/dev/null || echo "")
if echo "$HOSTNAME" | grep -qi "rescue"; then
    log "WARNING: Hostname is '$HOSTNAME' — server may still be in rescue mode!"
    exit 1
fi
log "Hostname: $HOSTNAME"

# Step 6: Verify key services
if [ ${#SERVICE_PORTS[@]} -gt 0 ]; then
    log "Checking service ports..."
    for port in "${SERVICE_PORTS[@]}"; do
        PORT_START=$(date +%s)
        while true; do
            PORT_ELAPSED=$(( $(date +%s) - PORT_START ))
            if [ "$PORT_ELAPSED" -ge "$MAX_SERVICE_WAIT" ]; then
                log "TIMEOUT: Port $port not responding after ${MAX_SERVICE_WAIT}s"
                exit 1
            fi
            if timeout 3 bash -c "echo > /dev/tcp/$HOST/$port" 2>/dev/null; then
                log "  ✅ Port $port: OPEN (after ${PORT_ELAPSED}s)"
                break
            fi
            sleep 2
        done
    done
fi

log "✅ Reboot complete — $HOST is up with all services"