#!/usr/bin/env bash

#####################################################################################
#### nginx-smoke-tests.sh — Verify Nginx is healthy and serving traffic          ####
#### Checks process, config, ports, vhosts, SSL, upstreams, errors, latency.     ####
####                                                                             ####
#### Author:  Phil Connor                                                        ####
#### Contact: contact@mylinux.work                                               ####
#### License: MIT                                                                ####
#### Version: 1.0                                                                ####
####                                                                             ####
#### Usage:   ./nginx-smoke-tests.sh                                             ####
####          VHOSTS=site.com:443,app.com:8080 ./nginx-smoke-tests.sh            ####
####                                                                             ####
#### See --help for all options.                                                 ####
#####################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
NGINX_HOST="${NGINX_HOST:-localhost}"
NGINX_PORT="${NGINX_PORT:-80}"
NGINX_SSL_PORT="${NGINX_SSL_PORT:-443}"
VHOSTS="${VHOSTS:-}"
UPSTREAMS="${UPSTREAMS:-}"
SSL_WARN_DAYS="${SSL_WARN_DAYS:-30}"
ERROR_LOG="${ERROR_LOG:-/var/log/nginx/error.log}"
STUB_STATUS_URL="${STUB_STATUS_URL:-}"
MAX_RESPONSE_MS="${MAX_RESPONSE_MS:-1000}"
EXPECTED_WORKERS="${EXPECTED_WORKERS:-}"
SKIP_SSL="${SKIP_SSL:-false}"
SKIP_VHOSTS="${SKIP_VHOSTS:-false}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
COLOR="${COLOR:-auto}"
VERBOSE="${VERBOSE:-false}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
SKIP=0
TOTAL=0
RESULTS=()
START_TIME=""

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1" detail="${2:-}"
    ((PASS++)) || true; ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name}${detail:+ (${detail})}"
    else echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_fail() {
    local name="$1" detail="${2:-}"
    ((FAIL++)) || true; ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    else echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_skip() {
    local name="$1" reason="${2:-}"
    ((SKIP++)) || true; ((TOTAL++)) || true
    RESULTS+=("SKIP|${name}|${reason}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name} # SKIP ${reason}"
    else echo -e "  ${YELLOW}⊘${RESET} ${name}${reason:+ — ${reason}}"; fi
}

# ── Helpers ───────────────────────────────────────────────────────────
has_cmd() { command -v "$1" >/dev/null 2>&1; }

section() {
    if [[ "$OUTPUT_FORMAT" != "tap" ]]; then echo ""; echo -e "${BOLD}$1${RESET}"; fi
}

# ── Cleanup ───────────────────────────────────────────────────────────
# shellcheck disable=SC2317
cleanup() {
    verbose "Cleanup complete."
}
trap cleanup EXIT

# ══════════════════════════════════════════════════════════════════════
#                          TEST FUNCTIONS
# ══════════════════════════════════════════════════════════════════════

# ── 1. Process running ───────────────────────────────────────────────
test_process_running() {
    if has_cmd systemctl; then
        if systemctl is-active --quiet nginx 2>/dev/null; then record_pass "Nginx process running" "systemctl active"
        else record_fail "Nginx process running" "systemctl inactive"; fi
    elif pgrep -x nginx >/dev/null 2>&1; then
        record_pass "Nginx process running" "pgrep found master"
    else
        record_fail "Nginx process running" "no nginx process found"
    fi
}

# ── 2. Config valid ──────────────────────────────────────────────────
test_config_valid() {
    if ! has_cmd nginx; then record_skip "Config syntax valid" "nginx binary not in PATH"; return; fi

    local output
    output=$(nginx -t 2>&1) || true
    if echo "$output" | grep -qi "syntax is ok"; then
        record_pass "Config syntax valid" "nginx -t passed"
    else
        local detail
        detail=$(echo "$output" | grep -i "emerg\|error" | head -1)
        record_fail "Config syntax valid" "${detail:-nginx -t failed}"
    fi
}

# ── 3. Listening ports ──────────────────────────────────────────────
test_listening_ports() {
    local check_tool=""
    if has_cmd ss; then check_tool="ss"
    elif has_cmd netstat; then check_tool="netstat"
    else record_skip "Listening on port ${NGINX_PORT}" "ss/netstat not available"; return; fi

    local listen_output
    if [[ "$check_tool" == "ss" ]]; then
        listen_output=$(ss -tlnp 2>/dev/null) || true
    else
        listen_output=$(netstat -tlnp 2>/dev/null) || true
    fi

    if echo "$listen_output" | grep -qE ":${NGINX_PORT}\b"; then
        record_pass "Listening on port ${NGINX_PORT}"
    else
        record_fail "Listening on port ${NGINX_PORT}" "port ${NGINX_PORT} not in listen state"
    fi

    if [[ "$SKIP_SSL" != "true" ]]; then
        if echo "$listen_output" | grep -qE ":${NGINX_SSL_PORT}\b"; then
            record_pass "Listening on port ${NGINX_SSL_PORT}"
        else
            record_fail "Listening on port ${NGINX_SSL_PORT}" "port ${NGINX_SSL_PORT} not in listen state"
        fi
    fi
}

# ── 4. HTTP response ────────────────────────────────────────────────
test_http_response() {
    if ! has_cmd curl; then record_skip "HTTP response" "curl not installed"; return; fi

    local http_code
    http_code=$(curl -sf -o /dev/null -w '%{http_code}' --max-time 10 \
        "http://${NGINX_HOST}:${NGINX_PORT}/" 2>/dev/null) || http_code="000"

    if [[ "$http_code" =~ ^(200|301|302|304)$ ]]; then
        record_pass "HTTP response" "HTTP ${http_code} from ${NGINX_HOST}:${NGINX_PORT}"
    else
        record_fail "HTTP response" "HTTP ${http_code} from ${NGINX_HOST}:${NGINX_PORT}"
    fi
}

# ── 5. Virtual host check ───────────────────────────────────────────
test_vhosts() {
    if [[ "$SKIP_VHOSTS" == "true" ]]; then record_skip "Virtual host check" "SKIP_VHOSTS=true"; return; fi
    if [[ -z "$VHOSTS" ]]; then record_skip "Virtual host check" "VHOSTS not set"; return; fi
    if ! has_cmd curl; then record_skip "Virtual host check" "curl not installed"; return; fi

    local IFS=','
    for entry in $VHOSTS; do
        local vhost="${entry%%:*}"
        local port="${entry##*:}"
        [[ "$vhost" == "$port" ]] && port="$NGINX_PORT"

        local scheme="http"
        [[ "$port" == "443" || "$port" == "$NGINX_SSL_PORT" ]] && scheme="https"

        local http_code
        http_code=$(curl -sf -o /dev/null -w '%{http_code}' --max-time 10 \
            -H "Host: ${vhost}" --insecure \
            "${scheme}://${NGINX_HOST}:${port}/" 2>/dev/null) || http_code="000"

        if [[ "$http_code" =~ ^(200|301|302|304)$ ]]; then
            record_pass "Vhost ${vhost}:${port}" "HTTP ${http_code}"
        else
            record_fail "Vhost ${vhost}:${port}" "HTTP ${http_code}"
        fi
    done
}

# ── 6. SSL certificate check ────────────────────────────────────────
test_ssl_cert() {
    if [[ "$SKIP_SSL" == "true" ]]; then record_skip "SSL certificate" "SKIP_SSL=true"; return; fi
    if ! has_cmd openssl; then record_skip "SSL certificate" "openssl not installed"; return; fi

    local hosts_to_check=("${NGINX_HOST}")
    if [[ -n "$VHOSTS" ]]; then
        local IFS=','
        for entry in $VHOSTS; do
            local vhost="${entry%%:*}"
            local port="${entry##*:}"
            [[ "$vhost" == "$port" ]] && port="$NGINX_PORT"
            [[ "$port" == "443" || "$port" == "$NGINX_SSL_PORT" ]] && hosts_to_check+=("$vhost")
        done
    fi

    local seen=()
    for host in "${hosts_to_check[@]}"; do
        # Deduplicate
        local already=false
        for s in "${seen[@]+"${seen[@]}"}"; do [[ "$s" == "$host" ]] && already=true; done
        [[ "$already" == "true" ]] && continue
        seen+=("$host")

        local expiry_str
        expiry_str=$(echo | openssl s_client -servername "$host" \
            -connect "${NGINX_HOST}:${NGINX_SSL_PORT}" 2>/dev/null \
            | openssl x509 -noout -enddate 2>/dev/null \
            | sed 's/notAfter=//') || true

        if [[ -z "$expiry_str" ]]; then
            record_fail "SSL cert ${host}" "could not retrieve certificate"
            continue
        fi

        local expiry_epoch now_epoch days_left
        expiry_epoch=$(date -d "$expiry_str" +%s 2>/dev/null) || true
        now_epoch=$(date +%s)

        if [[ -z "$expiry_epoch" ]]; then
            record_fail "SSL cert ${host}" "could not parse expiry date"
            continue
        fi

        days_left=$(( (expiry_epoch - now_epoch) / 86400 ))
        if [[ $days_left -lt 0 ]]; then
            record_fail "SSL cert ${host}" "expired ${days_left#-} days ago"
        elif [[ $days_left -lt $SSL_WARN_DAYS ]]; then
            record_fail "SSL cert ${host}" "expires in ${days_left}d (threshold ${SSL_WARN_DAYS}d)"
        else
            record_pass "SSL cert ${host}" "expires in ${days_left}d"
        fi
    done
}

# ── 7. Upstream health ──────────────────────────────────────────────
test_upstreams() {
    if [[ -z "$UPSTREAMS" ]]; then record_skip "Upstream health" "UPSTREAMS not set"; return; fi

    local IFS=','
    for entry in $UPSTREAMS; do
        local upstream_host="${entry%%:*}"
        local upstream_port="${entry##*:}"
        [[ "$upstream_host" == "$upstream_port" ]] && upstream_port="80"

        if has_cmd curl; then
            local http_code
            http_code=$(curl -sf -o /dev/null -w '%{http_code}' --max-time 5 \
                "http://${upstream_host}:${upstream_port}/" 2>/dev/null) || http_code="000"
            if [[ "$http_code" != "000" ]]; then
                record_pass "Upstream ${upstream_host}:${upstream_port}" "HTTP ${http_code}"
            else
                record_fail "Upstream ${upstream_host}:${upstream_port}" "no response"
            fi
        elif has_cmd nc; then
            if nc -z -w5 "$upstream_host" "$upstream_port" 2>/dev/null; then
                record_pass "Upstream ${upstream_host}:${upstream_port}" "port open"
            else
                record_fail "Upstream ${upstream_host}:${upstream_port}" "port closed"
            fi
        else
            record_skip "Upstream ${upstream_host}:${upstream_port}" "curl/nc not available"
        fi
    done
}

# ── 8. Error log check ──────────────────────────────────────────────
test_error_log() {
    if [[ ! -f "$ERROR_LOG" ]]; then
        record_skip "Error log check" "${ERROR_LOG} not found"
        return
    fi
    if [[ ! -r "$ERROR_LOG" ]]; then
        record_skip "Error log check" "${ERROR_LOG} not readable"
        return
    fi

    local critical_count
    critical_count=$(grep -ciE '\[(crit|alert|emerg)\]' "$ERROR_LOG" 2>/dev/null) || critical_count=0

    if [[ $critical_count -eq 0 ]]; then
        record_pass "Error log check" "no critical/alert/emerg entries"
    else
        record_fail "Error log check" "${critical_count} critical entries in ${ERROR_LOG}"
    fi
}

# ── 9. Worker count ─────────────────────────────────────────────────
test_worker_count() {
    local actual_workers
    actual_workers=$(pgrep -c 'nginx: worker' 2>/dev/null) || actual_workers=0

    if [[ $actual_workers -eq 0 ]]; then
        record_fail "Worker process count" "no worker processes found"
        return
    fi

    local expected="$EXPECTED_WORKERS"
    if [[ -z "$expected" ]]; then
        # Try auto-detect from config
        if has_cmd nginx; then
            local conf_val
            conf_val=$(nginx -T 2>/dev/null | grep -E '^\s*worker_processes' | awk '{print $2}' | tr -d ';' | head -1) || true
            if [[ "$conf_val" == "auto" ]]; then
                expected=$(nproc 2>/dev/null) || expected=""
            elif [[ "$conf_val" =~ ^[0-9]+$ ]]; then
                expected="$conf_val"
            fi
        fi
    fi

    if [[ -n "$expected" ]]; then
        if [[ $actual_workers -eq $expected ]]; then
            record_pass "Worker process count" "${actual_workers} workers (expected ${expected})"
        else
            record_fail "Worker process count" "${actual_workers} workers (expected ${expected})"
        fi
    else
        record_pass "Worker process count" "${actual_workers} workers"
    fi
}

# ── 10. Stub status ─────────────────────────────────────────────────
test_stub_status() {
    if [[ -z "$STUB_STATUS_URL" ]]; then record_skip "Stub status" "STUB_STATUS_URL not set"; return; fi
    if ! has_cmd curl; then record_skip "Stub status" "curl not installed"; return; fi

    local output
    output=$(curl -sf --max-time 5 "$STUB_STATUS_URL" 2>/dev/null) || true

    if echo "$output" | grep -q "Active connections"; then
        local active
        active=$(echo "$output" | grep "Active connections" | awk '{print $3}')
        record_pass "Stub status" "active connections: ${active}"
    else
        record_fail "Stub status" "no valid stub_status response from ${STUB_STATUS_URL}"
    fi
}

# ── 11. Response time ───────────────────────────────────────────────
test_response_time() {
    if ! has_cmd curl; then record_skip "Response time" "curl not installed"; return; fi

    local time_total
    time_total=$(curl -sf -o /dev/null -w '%{time_total}' --max-time 10 \
        "http://${NGINX_HOST}:${NGINX_PORT}/" 2>/dev/null) || true

    if [[ -z "$time_total" ]]; then
        record_fail "Response time" "request failed"
        return
    fi

    local ms
    ms=$(echo "$time_total" | awk '{printf "%.0f", $1 * 1000}') || ms=0

    if [[ $ms -le $MAX_RESPONSE_MS ]]; then
        record_pass "Response time" "${ms}ms (max ${MAX_RESPONSE_MS}ms)"
    else
        record_fail "Response time" "${ms}ms exceeds ${MAX_RESPONSE_MS}ms"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          OUTPUT
# ══════════════════════════════════════════════════════════════════════

print_tap_header() {
    echo "TAP version 13"
}

print_tap_footer() {
    echo "1..${TOTAL}"
    echo "# pass ${PASS}"
    echo "# fail ${FAIL}"
    echo "# skip ${SKIP}"
}

print_summary() {
    local end_time; end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))
    echo ""
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    echo -e "${BOLD}Summary${RESET}  Nginx Smoke Tests"
    echo -e "  ${GREEN}${PASS} passed${RESET}  ${RED}${FAIL} failed${RESET}  ${YELLOW}${SKIP} skipped${RESET}  (${duration}s)"
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    if [[ $FAIL -eq 0 ]]; then echo -e "${GREEN}${BOLD}All tests passed.${RESET}"
    else echo -e "${RED}${BOLD}${FAIL} test(s) failed.${RESET}"; fi
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Smoke-test Nginx process, config, ports, vhosts, SSL, upstreams, and latency.

Environment variables (all optional):
  NGINX_HOST          Target host                (default: localhost)
  NGINX_PORT          HTTP port                  (default: 80)
  NGINX_SSL_PORT      HTTPS port                 (default: 443)
  VHOSTS              Comma-separated host:port  (default: none)
  UPSTREAMS           Comma-separated host:port  (default: none)
  SSL_WARN_DAYS       Cert expiry threshold      (default: 30)
  ERROR_LOG           Error log path             (default: /var/log/nginx/error.log)
  STUB_STATUS_URL     stub_status endpoint URL   (default: none)
  MAX_RESPONSE_MS     Max response time in ms    (default: 1000)
  EXPECTED_WORKERS    Expected worker count      (default: auto-detect)
  SKIP_SSL            Skip SSL tests             (default: false)
  SKIP_VHOSTS         Skip vhost tests           (default: false)
  OUTPUT_FORMAT       text or tap                (default: text)
  COLOR               auto, always, never        (default: auto)
  VERBOSE             Show debug output          (default: false)

Options:
  --skip-ssl          Skip SSL certificate checks
  --skip-vhosts       Skip virtual host checks
  --format FORMAT     Output format: text (default), tap
  --verbose           Show debug output
  --no-color          Disable colored output
  --help              Show this help

Examples:
  ./$(basename "$0")
  VHOSTS=site.com:443,app.com:8080 ./$(basename "$0")
  UPSTREAMS=127.0.0.1:3000,127.0.0.1:3001 ./$(basename "$0")
  OUTPUT_FORMAT=tap ./$(basename "$0") --skip-ssl
EOF
}

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --skip-ssl)        SKIP_SSL=true ;;
            --skip-vhosts)     SKIP_VHOSTS=true ;;
            --format)          OUTPUT_FORMAT="$2"; shift ;;
            --verbose)         VERBOSE=true ;;
            --no-color)        COLOR=never ;;
            --help|-h)         usage; exit 0 ;;
            *)                 err "Unknown option: $1"; usage; exit 1 ;;
        esac
        shift
    done

    setup_colors

    if [[ "${BASH_VERSINFO[0]}" -lt 4 ]]; then
        err "Bash 4+ is required"
        exit 1
    fi

    START_TIME=$(date +%s)

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_header
    else
        echo ""
        echo -e "${BOLD}Nginx Smoke Tests${RESET}"
        echo -e "Host: ${NGINX_HOST}  HTTP: ${NGINX_PORT}  HTTPS: ${NGINX_SSL_PORT}"
        echo -e "Time: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    fi

    section "Process & Config"
    test_process_running
    test_config_valid
    test_worker_count

    section "Network"
    test_listening_ports
    test_http_response
    test_response_time

    section "Virtual Hosts"
    test_vhosts

    section "SSL"
    test_ssl_cert

    section "Upstreams"
    test_upstreams

    section "Health"
    test_stub_status
    test_error_log

    # ── Results ──
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_footer
    else
        print_summary
    fi

    [[ $FAIL -eq 0 ]] && exit 0 || exit 1
}

main "$@"
