#!/bin/sh
# rolling-upgrade.sh — Zero-downtime upgrade for a running fafb-q server.
#
# Orchestrates: preflight → drain → shutdown wait → binary swap → restart → health check.
#
# Usage:
#   FAFB_ADMIN_API_KEY=mykey ./scripts/rolling-upgrade.sh
#   FAFB_ADMIN_API_KEY=mykey ./scripts/rolling-upgrade.sh --binary /path/to/fafb-q
#   ./scripts/rolling-upgrade.sh --admin-url http://host:8080 --admin-key mykey
#   ./scripts/rolling-upgrade.sh --force
#   ./scripts/rolling-upgrade.sh --config /etc/fafb-q/config.toml
#
# Prerequisites:
#   - A running fafb-q server with admin API enabled
#   - FAFB_ADMIN_API_KEY set (env var or --admin-key flag)
#   - curl installed
#
# The script uses the admin drain API (POST /admin/drain) to trigger a graceful
# shutdown. The server drains connections, sends GoingAway to capable clients,
# and exits on its own. The script then installs the new binary and restarts.
#
# For the full deploy lifecycle, see: docs/deploy.md
set -eu

# ── Constants ──────────────────────────────────────────────────────────────────

BINARY_NAME="fafb-q"
DEFAULT_INSTALL_DIR="/usr/local/bin"
BASE_URL="https://bitbucket.org/thermsio/fafb-q/downloads"

# ── Helpers ────────────────────────────────────────────────────────────────────

# Run a command with sudo if not root. Same pattern as install.sh.
needs_sudo=false
run() {
    if [ "$needs_sudo" = true ]; then
        sudo "$@"
    else
        "$@"
    fi
}

die() {
    echo "Error: $1" >&2
    exit 1
}

usage() {
    cat <<'USAGE'
Usage: rolling-upgrade.sh [OPTIONS]

Orchestrate a zero-downtime upgrade of a running fafb-q server.

Options:
  --admin-url URL       Admin endpoint base URL (default: http://localhost:8080)
  --admin-key KEY       Admin API key (default: $FAFB_ADMIN_API_KEY env var)
  --binary PATH         Path to pre-staged binary (default: download latest)
  --force               SIGKILL if drain/shutdown times out
  --drain-timeout SECS  Seconds to wait for shutdown (default: 30)
  --config PATH         Config file for server restart (required in non-systemd mode)
  --help                Show this help message

Environment:
  FAFB_ADMIN_API_KEY    Admin API key (alternative to --admin-key)

Examples:
  FAFB_ADMIN_API_KEY=mykey ./scripts/rolling-upgrade.sh
  ./scripts/rolling-upgrade.sh --binary /tmp/fafb-q --admin-key mykey
  ./scripts/rolling-upgrade.sh --force --drain-timeout 60

Full deploy documentation: docs/deploy.md
USAGE
    exit 0
}

# ── Flag Parsing ───────────────────────────────────────────────────────────────
# Each flag has a default. --admin-key falls back to $FAFB_ADMIN_API_KEY env var.
# --config is only required in non-systemd (direct PID) mode for restart.

ADMIN_URL="http://localhost:8080"
ADMIN_KEY="${FAFB_ADMIN_API_KEY:-}"
BINARY_PATH=""
FORCE=false
DRAIN_TIMEOUT=30
CONFIG_PATH=""

while [ $# -gt 0 ]; do
    case "$1" in
        --admin-url)  ADMIN_URL="$2";  shift 2 ;;
        --admin-key)  ADMIN_KEY="$2";  shift 2 ;;
        --binary)     BINARY_PATH="$2"; shift 2 ;;
        --force)      FORCE=true;       shift ;;
        --drain-timeout)
            case "$2" in
                ''|*[!0-9]*) die "--drain-timeout requires a numeric value, got: $2" ;;
            esac
            DRAIN_TIMEOUT="$2"; shift 2 ;;
        --config)     CONFIG_PATH="$2"; shift 2 ;;
        --help)       usage ;;
        *)            die "unknown flag: $1. Use --help for usage." ;;
    esac
done

# ── Main ───────────────────────────────────────────────────────────────────────

main() {
    # ── Step 1: Preflight Checks ──────────────────────────────────────────
    # Validate everything before touching the running server.
    # If anything fails here, we exit without side effects.

    if ! command -v curl >/dev/null 2>&1; then
        die "curl is required but not found"
    fi

    if [ -z "$ADMIN_KEY" ]; then
        die "admin key required. Set FAFB_ADMIN_API_KEY or use --admin-key."
    fi

    # Validate the new binary first (pure input validation — no network needed).
    # This catches obvious errors (bad --binary path) before we talk to the server.
    tmpdir="$(mktemp -d)"
    trap 'rm -rf "$tmpdir"' EXIT INT TERM

    if [ -n "$BINARY_PATH" ]; then
        # Pre-staged binary provided
        if [ ! -x "$BINARY_PATH" ]; then
            die "binary not found or not executable: $BINARY_PATH"
        fi
        new_binary="$BINARY_PATH"
    else
        # Download latest — same logic as install.sh
        os="$(uname -s)"
        if [ "$os" != "Linux" ]; then
            die "unsupported OS: $os. This script only supports Linux."
        fi

        arch="$(uname -m)"
        case "$arch" in
            x86_64)        target="x86_64-unknown-linux-musl" ;;
            aarch64|arm64) target="aarch64-unknown-linux-musl" ;;
            *)             die "unsupported architecture: $arch" ;;
        esac

        tarball="fafb-q-latest-${target}.tar.gz"
        echo "Preflight: downloading fafb-q for ${target}..."
        curl -fsSL --connect-timeout 10 --max-time 120 "${BASE_URL}/${tarball}" -o "${tmpdir}/${tarball}"
        tar xzf "${tmpdir}/${tarball}" --strip-components=1 -C "$tmpdir"

        if [ ! -x "${tmpdir}/${BINARY_NAME}" ]; then
            die "binary not found or not executable after extraction"
        fi
        new_binary="${tmpdir}/${BINARY_NAME}"
    fi

    # Smoke test the new binary
    new_version=$("$new_binary" --version 2>&1) || die "new binary failed to execute"
    echo "Preflight: new binary version: $new_version"

    # Verify server is reachable and healthy (after binary validation —
    # this is the last preflight check before we start affecting the server).
    echo "Preflight: checking server at $ADMIN_URL..."
    http_code=$(curl -s -o /dev/null -w '%{http_code}' --connect-timeout 5 --max-time 10 "$ADMIN_URL/health/ready") || true
    if [ "$http_code" != "200" ]; then
        die "server not reachable or not ready at $ADMIN_URL/health/ready (got HTTP $http_code)"
    fi
    echo "Preflight: server reachable at $ADMIN_URL"

    # ── Step 2: Resolve Process Management ────────────────────────────────
    # Detect whether the server is managed by systemd or running as a bare
    # process. This determines how we wait for shutdown and restart.
    # - systemd: the drain API triggers exit; we poll systemctl is-active.
    # - direct PID: we find the PID and poll kill -0 until it's gone.

    use_systemd=false
    server_pid=""
    install_path="$DEFAULT_INSTALL_DIR/$BINARY_NAME"

    if command -v systemctl >/dev/null 2>&1 && systemctl list-unit-files fafb-q.service >/dev/null 2>&1; then
        use_systemd=true
        echo "Detected: systemd service (fafb-q.service)"
    else
        # Direct PID mode — find the running process.
        # pidof/pgrep may return multiple PIDs; take only the first one.
        # If multiple instances are running, the operator should use systemd
        # or specify which instance to upgrade.
        if command -v pidof >/dev/null 2>&1; then
            server_pid=$(pidof "$BINARY_NAME" 2>/dev/null | awk '{print $1}') || true
        fi
        if [ -z "$server_pid" ] && command -v pgrep >/dev/null 2>&1; then
            server_pid=$(pgrep -x "$BINARY_NAME" 2>/dev/null | head -1) || true
        fi
        if [ -z "$server_pid" ]; then
            die "no running fafb-q process found"
        fi
        echo "Detected: direct PID mode (PID $server_pid)"

        # Resolve the binary path from /proc for install target.
        # Requires same user or root. Falls back to default if unreadable.
        if [ -L "/proc/$server_pid/exe" ]; then
            resolved=$(readlink "/proc/$server_pid/exe" 2>/dev/null) || true
            if [ -n "$resolved" ]; then
                install_path="$resolved"
            fi
        fi
    fi

    # ── Step 3: Trigger Drain ─────────────────────────────────────────────
    # POST /admin/drain triggers the server's graceful shutdown sequence:
    #   force-flush → Draining → GoingAway to clients → wait → ShuttingDown → exit
    # The server exits on its own after drain completes. We do NOT need to
    # send SIGTERM or call systemctl stop in the normal path.

    echo "Draining server..."
    drain_code=$(curl -s -o /dev/null -w '%{http_code}' --connect-timeout 5 --max-time 10 -X POST "$ADMIN_URL/admin/drain" \
        -H "Authorization: Bearer $ADMIN_KEY" \
        -H "Content-Type: application/json" \
        -d '{"reason": "upgrade"}') || true

    case "$drain_code" in
        200)
            echo "Drain accepted. Waiting for server to shut down..."
            ;;
        409)
            # Server is already draining (e.g., another operator triggered it).
            # This is not an error — proceed to wait for shutdown.
            echo "Warning: server is already draining. Proceeding to wait for shutdown."
            ;;
        *)
            die "drain request failed (HTTP $drain_code). Check admin key and URL."
            ;;
    esac

    # Poll /health/ready until it returns 503 (confirms drain mode entered).
    # Timeout after 10s — the drain was accepted; ready may already be 503.
    drain_confirm_waited=0
    while [ "$drain_confirm_waited" -lt 10 ]; do
        ready_code=$(curl -s -o /dev/null -w '%{http_code}' --connect-timeout 2 --max-time 5 "$ADMIN_URL/health/ready") || true
        if [ "$ready_code" = "503" ] || [ "$ready_code" = "000" ]; then
            break
        fi
        sleep 2
        drain_confirm_waited=$((drain_confirm_waited + 2))
    done

    # ── Step 4: Wait for Shutdown ─────────────────────────────────────────
    # The server should exit on its own after drain completes. We poll until
    # the process is gone. If it doesn't exit within --drain-timeout, we
    # either abort (safe default) or SIGKILL (--force).

    waited=0
    while [ "$waited" -lt "$DRAIN_TIMEOUT" ]; do
        if [ "$use_systemd" = true ]; then
            status=$(systemctl is-active fafb-q 2>/dev/null) || true
            if [ "$status" = "inactive" ] || [ "$status" = "failed" ]; then
                break
            fi
        else
            if ! kill -0 "$server_pid" 2>/dev/null; then
                break
            fi
        fi
        sleep 2
        waited=$((waited + 2))
    done

    # Check if server is still running
    still_running=false
    if [ "$use_systemd" = true ]; then
        status=$(systemctl is-active fafb-q 2>/dev/null) || true
        if [ "$status" = "active" ] || [ "$status" = "activating" ]; then
            still_running=true
        fi
    else
        if kill -0 "$server_pid" 2>/dev/null; then
            still_running=true
        fi
    fi

    if [ "$still_running" = true ]; then
        if [ "$FORCE" = true ]; then
            echo "Drain timed out. Force-killing server..."
            if [ "$use_systemd" = true ]; then
                systemctl kill --signal=KILL fafb-q 2>/dev/null || true
            else
                kill -9 "$server_pid" 2>/dev/null || true
            fi
            # Wait up to 5 seconds for the process to die
            force_waited=0
            while [ "$force_waited" -lt 5 ]; do
                if [ "$use_systemd" = true ]; then
                    status=$(systemctl is-active fafb-q 2>/dev/null) || true
                    if [ "$status" = "inactive" ] || [ "$status" = "failed" ]; then
                        break
                    fi
                else
                    if ! kill -0 "$server_pid" 2>/dev/null; then
                        break
                    fi
                fi
                sleep 1
                force_waited=$((force_waited + 1))
            done
            # Final check
            if [ "$use_systemd" = true ]; then
                status=$(systemctl is-active fafb-q 2>/dev/null) || true
                if [ "$status" = "active" ]; then
                    die "failed to kill server even with --force"
                fi
            else
                if kill -0 "$server_pid" 2>/dev/null; then
                    die "failed to kill PID $server_pid even with --force"
                fi
            fi
        else
            die "drain timed out after ${DRAIN_TIMEOUT}s. Server is still running. Use --force to kill, or increase --drain-timeout."
        fi
    fi

    if [ -n "$server_pid" ]; then
        echo "Server stopped (was PID $server_pid)."
    else
        echo "Server stopped."
    fi

    # ── Step 5: Install Binary ────────────────────────────────────────────
    # Install the new binary to the same path as the old one.
    # Uses sudo if not root (same pattern as install.sh).

    if [ "$(id -u)" -ne 0 ]; then
        if command -v sudo >/dev/null 2>&1; then
            needs_sudo=true
        else
            die "permission denied. Run as root or install sudo."
        fi
    fi

    echo "Installing $new_version to $install_path..."
    run mkdir -p "$(dirname "$install_path")"
    run install -m 755 "$new_binary" "$install_path"

    # ── Step 6: Start and Verify ──────────────────────────────────────────
    # Restart the server and verify it becomes healthy.
    # - systemd: systemctl start, check is-active, poll /health/ready.
    # - direct PID: nohup with --config, poll /health/ready.
    # - no --config in PID mode: skip start, tell user to start manually.

    started=false
    if [ "$use_systemd" = true ]; then
        echo "Starting server via systemctl..."
        systemctl start fafb-q

        # Verify systemd reports the service as active
        sleep 1
        status=$(systemctl is-active fafb-q 2>/dev/null) || true
        if [ "$status" = "failed" ]; then
            echo "Error: server failed to start. Check logs:" >&2
            systemctl status fafb-q --no-pager 2>&1 || true
            die "server failed to start. Check logs with: journalctl -u fafb-q"
        fi
        started=true
    elif [ -n "$CONFIG_PATH" ]; then
        echo "Starting server..."
        nohup "$install_path" serve --config "$CONFIG_PATH" > /dev/null 2>&1 &
        new_pid=$!
        echo "Server started (PID $new_pid)."
        started=true
    else
        echo "Binary installed. Start the server manually."
    fi

    # Poll /health/ready until it returns 200 (server initialized).
    # Timeout: 60s — accounts for large dataset replay during engine.init().
    if [ "$started" = true ]; then
        echo "Waiting for server to become ready..."
        health_waited=0
        while [ "$health_waited" -lt 60 ]; do
            ready_code=$(curl -s -o /dev/null -w '%{http_code}' --connect-timeout 2 --max-time 5 "$ADMIN_URL/health/ready") || true
            if [ "$ready_code" = "200" ]; then
                break
            fi
            sleep 2
            health_waited=$((health_waited + 2))
        done

        if [ "$ready_code" = "200" ]; then
            echo "Health check: /health/ready returned 200"
            "$install_path" --version 2>/dev/null || true
            echo "Upgrade complete."
        else
            echo "Warning: server started but /health/ready not yet 200 after 60s. It may still be initializing." >&2
            exit 1
        fi
    fi
}

main
