#!/usr/bin/env bash
# ============================================================================
# RevSocks v4 — Watchdog Service
#
# Monitors the RevSocks server and CF bridge, restarts on failure,
# rotates logs, checks system resources, and sends webhook alerts.
#
# Installation:
#   chmod +x watchdog.sh
#   # Run directly:
#   ./watchdog.sh
#   # Or via cron (every 2 minutes):
#   */2 * * * * /root/revsocks_v4/watchdog.sh >> /var/log/revsocks_watchdog.log 2>&1
#
# Configuration (environment or edit below):
#   REVSOCKS_DIR          — Project root (auto-detected)
#   REVSOCKS_WEBHOOK_URL  — Webhook URL for alerts (Slack/Discord/Teams)
#   REVSOCKS_LOG_MAX_MB   — Max log size before rotation (default: 1)
#   REVSOCKS_SYSTEMD_SVC  — Systemd service name (default: revsocks)
#   REVSOCKS_TMUX_SESSION — Tmux session name (default: revsocks)
#
# ============================================================================
set -uo pipefail

# --- Configuration ---
REVSOCKS_DIR="${REVSOCKS_DIR:-$(cd "$(dirname "$0")" && pwd)}"
WEBHOOK_URL="${REVSOCKS_WEBHOOK_URL:-}"
LOG_MAX_MB="${REVSOCKS_LOG_MAX_MB:-1}"
SYSTEMD_SVC="${REVSOCKS_SYSTEMD_SVC:-revsocks}"
TMUX_SESSION="${REVSOCKS_TMUX_SESSION:-revsocks}"
SERVER_SCRIPT="${REVSOCKS_DIR}/server/server.py"
BRIDGE_SCRIPT="${REVSOCKS_DIR}/server/cf_bridge.py"
LOG_DIR="${REVSOCKS_DIR}/output/logs"
WATCHDOG_LOG="${LOG_DIR}/watchdog.log"
SERVER_LOG="${LOG_DIR}/server.log"
BRIDGE_LOG="${LOG_DIR}/bridge.log"
LISTEN_PORT=443
BRIDGE_PORT=4443
MIN_DISK_MB=500
MIN_MEM_MB=128

# --- Ensure log directory ---
mkdir -p "$LOG_DIR"

# --- Timestamp ---
ts() { date '+%Y-%m-%d %H:%M:%S'; }

# --- Logging ---
log_info()  { echo "[$(ts)] [INFO]  $*" | tee -a "$WATCHDOG_LOG"; }
log_warn()  { echo "[$(ts)] [WARN]  $*" | tee -a "$WATCHDOG_LOG"; }
log_error() { echo "[$(ts)] [ERROR] $*" | tee -a "$WATCHDOG_LOG"; }
log_ok()    { echo "[$(ts)] [OK]    $*" | tee -a "$WATCHDOG_LOG"; }

# --- Webhook alert ---
send_alert() {
    local message="$1"
    local hostname
    hostname=$(hostname 2>/dev/null || echo "unknown")

    if [ -n "$WEBHOOK_URL" ]; then
        local payload
        # Detect webhook type from URL and format accordingly
        if echo "$WEBHOOK_URL" | grep -q "discord"; then
            payload="{\"content\":\"**[RevSocks Watchdog]** (${hostname})\\n${message}\"}"
        elif echo "$WEBHOOK_URL" | grep -q "slack"; then
            payload="{\"text\":\"*[RevSocks Watchdog]* (${hostname})\\n${message}\"}"
        else
            # Generic JSON webhook (Teams, custom)
            payload="{\"title\":\"RevSocks Watchdog\",\"host\":\"${hostname}\",\"message\":\"${message}\",\"timestamp\":\"$(ts)\"}"
        fi

        curl -s -m 10 -X POST \
             -H "Content-Type: application/json" \
             -d "$payload" \
             "$WEBHOOK_URL" >/dev/null 2>&1 || log_warn "Webhook delivery failed"
    fi
}

# --- Check if process is running ---
is_running() {
    pgrep -f "$1" >/dev/null 2>&1
}

# --- Check if port is listening ---
is_port_open() {
    local port="$1"
    if command -v ss &>/dev/null; then
        ss -tlnp 2>/dev/null | grep -q ":${port} "
    elif command -v netstat &>/dev/null; then
        netstat -tlnp 2>/dev/null | grep -q ":${port} "
    else
        (echo >/dev/tcp/127.0.0.1/"$port") 2>/dev/null
    fi
}

# --- Restart server ---
restart_server() {
    local reason="$1"
    log_warn "Restarting server: $reason"
    send_alert "Restarting server: ${reason}"

    # Method 1: Systemd (preferred)
    if systemctl is-active --quiet "$SYSTEMD_SVC" 2>/dev/null || \
       systemctl list-unit-files 2>/dev/null | grep -q "^${SYSTEMD_SVC}"; then
        log_info "Restarting via systemd: $SYSTEMD_SVC"
        systemctl restart "$SYSTEMD_SVC" 2>/dev/null
        sleep 3
        if systemctl is-active --quiet "$SYSTEMD_SVC" 2>/dev/null; then
            log_ok "Server restarted via systemd"
            return 0
        else
            log_error "Systemd restart failed"
        fi
    fi

    # Method 2: Tmux
    if command -v tmux &>/dev/null; then
        # Kill old process
        pkill -f "python3.*server\.py" 2>/dev/null || true
        sleep 1

        # Check if tmux session exists
        if tmux has-session -t "$TMUX_SESSION" 2>/dev/null; then
            tmux send-keys -t "$TMUX_SESSION" C-c 2>/dev/null || true
            sleep 1
            tmux send-keys -t "$TMUX_SESSION" "python3 $SERVER_SCRIPT >> $SERVER_LOG 2>&1 &" Enter
        else
            tmux new-session -d -s "$TMUX_SESSION" \
                "python3 $SERVER_SCRIPT >> $SERVER_LOG 2>&1; exec bash"
        fi
        sleep 3

        if is_running "server\.py"; then
            log_ok "Server restarted via tmux"
            return 0
        else
            log_error "Tmux restart failed"
        fi
    fi

    # Method 3: Direct background
    pkill -f "python3.*server\.py" 2>/dev/null || true
    sleep 1
    nohup python3 "$SERVER_SCRIPT" >> "$SERVER_LOG" 2>&1 &
    sleep 3

    if is_running "server\.py"; then
        log_ok "Server restarted (background)"
        return 0
    fi

    log_error "All restart methods failed"
    send_alert "CRITICAL: All server restart methods failed"
    return 1
}

# --- Restart CF bridge ---
restart_bridge() {
    local reason="$1"
    log_warn "Restarting CF bridge: $reason"

    pkill -f "python3.*cf_bridge\.py" 2>/dev/null || true
    sleep 1
    nohup python3 "$BRIDGE_SCRIPT" >> "$BRIDGE_LOG" 2>&1 &
    sleep 2

    if is_running "cf_bridge\.py"; then
        log_ok "CF bridge restarted"
    else
        log_error "CF bridge restart failed"
        send_alert "CF bridge restart failed"
    fi
}

# --- Rotate log if > threshold ---
rotate_log() {
    local logfile="$1"
    local max_bytes=$((LOG_MAX_MB * 1024 * 1024))

    if [ -f "$logfile" ]; then
        local size
        size=$(stat -c%s "$logfile" 2>/dev/null || stat -f%z "$logfile" 2>/dev/null || echo 0)
        if [ "$size" -gt "$max_bytes" ]; then
            local rotated="${logfile}.$(date +%Y%m%d_%H%M%S).old"
            mv "$logfile" "$rotated"
            # Compress old log
            gzip "$rotated" 2>/dev/null &
            log_info "Rotated $logfile ($((size / 1024)) KB)"

            # Keep only last 5 rotated logs
            ls -t "${logfile}".*.old.gz 2>/dev/null | tail -n +6 | xargs rm -f 2>/dev/null || true
        fi
    fi
}

# --- Check disk space ---
check_disk() {
    local avail_kb
    avail_kb=$(df -k "$REVSOCKS_DIR" 2>/dev/null | tail -1 | awk '{print $4}')
    if [ -n "$avail_kb" ]; then
        local avail_mb=$((avail_kb / 1024))
        if [ "$avail_mb" -lt "$MIN_DISK_MB" ]; then
            log_error "Low disk space: ${avail_mb}MB available (min: ${MIN_DISK_MB}MB)"
            send_alert "Low disk space: ${avail_mb}MB available"
            return 1
        fi
    fi
    return 0
}

# --- Check memory ---
check_memory() {
    if [ -f /proc/meminfo ]; then
        local avail_kb
        avail_kb=$(grep MemAvailable /proc/meminfo 2>/dev/null | awk '{print $2}')
        if [ -n "$avail_kb" ]; then
            local avail_mb=$((avail_kb / 1024))
            if [ "$avail_mb" -lt "$MIN_MEM_MB" ]; then
                log_error "Low memory: ${avail_mb}MB available (min: ${MIN_MEM_MB}MB)"
                send_alert "Low memory: ${avail_mb}MB available"
                return 1
            fi
        fi
    fi
    return 0
}

# ============================================================================
# Main watchdog logic
# ============================================================================

log_info "--- Watchdog check started ---"
errors=0

# 1. Rotate logs
rotate_log "$SERVER_LOG"
rotate_log "$BRIDGE_LOG"
rotate_log "$WATCHDOG_LOG"

# 2. Check disk space
if ! check_disk; then
    errors=$((errors + 1))
fi

# 3. Check memory
if ! check_memory; then
    errors=$((errors + 1))
fi

# 4. Check server process
if ! is_running "server\.py"; then
    log_error "Server process not running"
    errors=$((errors + 1))
    if [ -f "$SERVER_SCRIPT" ]; then
        restart_server "process not found"
    else
        log_error "Server script not found: $SERVER_SCRIPT"
        send_alert "Server script missing: $SERVER_SCRIPT"
    fi
elif ! is_port_open "$LISTEN_PORT"; then
    log_error "Server running but port $LISTEN_PORT not listening"
    errors=$((errors + 1))
    restart_server "port $LISTEN_PORT not listening"
else
    log_ok "Server OK (process running, port $LISTEN_PORT listening)"
fi

# 5. Check CF bridge (optional — only if script exists)
if [ -f "$BRIDGE_SCRIPT" ]; then
    if ! is_running "cf_bridge\.py"; then
        log_warn "CF bridge not running"
        restart_bridge "process not found"
    elif ! is_port_open "$BRIDGE_PORT"; then
        log_warn "CF bridge running but port $BRIDGE_PORT not listening"
        restart_bridge "port $BRIDGE_PORT not listening"
    else
        log_ok "CF bridge OK (port $BRIDGE_PORT listening)"
    fi
fi

# 6. Summary
if [ "$errors" -gt 0 ]; then
    log_warn "Watchdog found $errors issue(s)"
else
    log_ok "All checks passed"
fi

log_info "--- Watchdog check complete ---"
exit 0