#!/usr/bin/env bash

#####################################################################################
#### database-smoke-tests.sh — Verify database health                            ####
#### Checks connectivity, auth, replication, backup age, bloat, connections.     ####
#### Supports: PostgreSQL, MySQL/MariaDB, Redis                                  ####
####                                                                             ####
#### Author:  Phil Connor                                                        ####
#### Contact: contact@mylinux.work                                               ####
#### License: MIT                                                                ####
#### Version: 1.0                                                                ####
####                                                                             ####
#### Usage:   DB_TYPE=postgresql DB_HOST=localhost ./database-smoke-tests.sh      ####
####          DB_TYPE=redis REDIS_HOST=localhost ./database-smoke-tests.sh        ####
####                                                                             ####
#### See --help for all options.                                                 ####
#####################################################################################

set -euo pipefail

# ---------------------------------------------------------------------------
# Help
# ---------------------------------------------------------------------------
show_help() {
    cat <<'EOF'
database-smoke-tests.sh — Database health smoke testing

ENVIRONMENT VARIABLES:
  DB_TYPE              (required) postgresql | mysql | redis
  DB_HOST              Database host (default: localhost)
  DB_PORT              Database port (default: auto — 5432/3306/6379)
  DB_USER              Database user (default: postgres | root | "")
  DB_PASS              Database password (default: "")
  DB_NAME              Database name  (default: postgres | mysql)

  REDIS_HOST           Redis host  (falls back to DB_HOST)
  REDIS_PORT           Redis port  (falls back to DB_PORT)
  REDIS_AUTH           Redis auth  (falls back to DB_PASS)

  MAX_REPLICATION_LAG_S  Max replication lag in seconds (default: 30)
  MAX_BACKUP_AGE_H       Max backup / last-save age in hours (default: 26)
  MAX_CONNECTIONS_PCT    Connection usage threshold %     (default: 80)
  SKIP_REPLICATION       Skip replication checks (default: false)
  SKIP_BACKUP_AGE        Skip backup-age checks  (default: false)

  OUTPUT_FORMAT        text | tap  (default: text)
  COLOR                auto | always | never (default: auto)
  VERBOSE              true | false (default: false)

EXAMPLES:
  DB_TYPE=postgresql DB_HOST=db1 DB_PASS=secret ./database-smoke-tests.sh
  DB_TYPE=mysql DB_HOST=db2 DB_USER=app DB_NAME=mydb ./database-smoke-tests.sh
  DB_TYPE=redis REDIS_HOST=cache1 REDIS_AUTH=pass ./database-smoke-tests.sh
EOF
    exit 0
}

[[ "${1:-}" == "--help" || "${1:-}" == "-h" ]] && show_help

# ---------------------------------------------------------------------------
# Environment defaults
# ---------------------------------------------------------------------------
DB_TYPE="${DB_TYPE:-}"
DB_HOST="${DB_HOST:-localhost}"
DB_PORT="${DB_PORT:-}"
DB_USER="${DB_USER:-}"
DB_PASS="${DB_PASS:-}"
DB_NAME="${DB_NAME:-}"

REDIS_HOST="${REDIS_HOST:-$DB_HOST}"
REDIS_PORT="${REDIS_PORT:-${DB_PORT:-6379}}"
REDIS_AUTH="${REDIS_AUTH:-$DB_PASS}"

MAX_REPLICATION_LAG_S="${MAX_REPLICATION_LAG_S:-30}"
MAX_BACKUP_AGE_H="${MAX_BACKUP_AGE_H:-26}"
MAX_CONNECTIONS_PCT="${MAX_CONNECTIONS_PCT:-80}"
SKIP_REPLICATION="${SKIP_REPLICATION:-false}"
SKIP_BACKUP_AGE="${SKIP_BACKUP_AGE:-false}"

OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
COLOR="${COLOR:-auto}"
VERBOSE="${VERBOSE:-false}"

# ---------------------------------------------------------------------------
# Apply per-engine defaults after DB_TYPE is known
# ---------------------------------------------------------------------------
apply_defaults() {
    case "$DB_TYPE" in
        postgresql)
            DB_PORT="${DB_PORT:-5432}"
            DB_USER="${DB_USER:-postgres}"
            DB_NAME="${DB_NAME:-postgres}"
            ;;
        mysql)
            DB_PORT="${DB_PORT:-3306}"
            DB_USER="${DB_USER:-root}"
            DB_NAME="${DB_NAME:-mysql}"
            ;;
        redis)
            REDIS_PORT="${REDIS_PORT:-6379}"
            ;;
        *)
            echo "ERROR: DB_TYPE must be one of: postgresql, mysql, redis" >&2
            exit 1
            ;;
    esac
}

# ---------------------------------------------------------------------------
# Colour setup
# ---------------------------------------------------------------------------
setup_colors() {
    RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    local use_color=false
    case "$COLOR" in
        always) use_color=true ;;
        never)  use_color=false ;;
        auto)   [[ -t 1 ]] && use_color=true ;;
    esac
    if $use_color; then
        RED=$'\033[0;31m'
        GREEN=$'\033[0;32m'
        YELLOW=$'\033[1;33m'
        BLUE=$'\033[0;34m'
        BOLD=$'\033[1m'
        RESET=$'\033[0m'
    fi
}

# ---------------------------------------------------------------------------
# Counters
# ---------------------------------------------------------------------------
PASS_COUNT=0
FAIL_COUNT=0
SKIP_COUNT=0
TEST_NUM=0

# ---------------------------------------------------------------------------
# run_test  "description"  command...
# ---------------------------------------------------------------------------
run_test() {
    local desc="$1"; shift
    TEST_NUM=$((TEST_NUM + 1))
    local output rc
    output=$("$@" 2>&1) && rc=0 || rc=$?

    if [[ $rc -eq 0 ]]; then
        PASS_COUNT=$((PASS_COUNT + 1))
        if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
            echo "ok $TEST_NUM - $desc"
        else
            echo "  ${GREEN}PASS${RESET}  $desc"
        fi
    elif [[ $rc -eq 2 ]]; then
        SKIP_COUNT=$((SKIP_COUNT + 1))
        if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
            echo "ok $TEST_NUM - $desc # SKIP ${output:-skipped}"
        else
            echo "  ${YELLOW}SKIP${RESET}  $desc — ${output:-skipped}"
        fi
    else
        FAIL_COUNT=$((FAIL_COUNT + 1))
        if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
            echo "not ok $TEST_NUM - $desc"
            [[ -n "$output" ]] && echo "# $output"
        else
            echo "  ${RED}FAIL${RESET}  $desc"
            [[ -n "$output" ]] && echo "        $output"
        fi
    fi

    if [[ "$VERBOSE" == "true" && -n "$output" && $rc -eq 0 ]]; then
        echo "        ${BLUE}→${RESET} $output"
    fi
}

# ---------------------------------------------------------------------------
# skip_test  "description"  "reason"
# ---------------------------------------------------------------------------
skip_test() {
    local desc="$1" reason="${2:-skipped}"
    TEST_NUM=$((TEST_NUM + 1))
    SKIP_COUNT=$((SKIP_COUNT + 1))
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok $TEST_NUM - $desc # SKIP $reason"
    else
        echo "  ${YELLOW}SKIP${RESET}  $desc — $reason"
    fi
}

# ---------------------------------------------------------------------------
# check_port  host  port  [timeout]
# ---------------------------------------------------------------------------
check_port() {
    local host="$1" port="$2" timeout="${3:-5}"
    if command -v nc &>/dev/null; then
        nc -z -w "$timeout" "$host" "$port" 2>/dev/null
    elif [[ -e /dev/tcp ]]; then
        timeout "$timeout" bash -c "echo >/dev/tcp/$host/$port" 2>/dev/null
    else
        (echo >/dev/tcp/"$host"/"$port") 2>/dev/null
    fi
}

# ---------------------------------------------------------------------------
# Helper: build psql / mysql invocations
# ---------------------------------------------------------------------------
run_psql() {
    PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" \
        -d "${1:-$DB_NAME}" -t -A -c "$2" 2>&1
}

run_mysql() {
    MYSQL_PWD="$DB_PASS" mysql -h "$DB_HOST" -P "$DB_PORT" -u "$DB_USER" \
        -D "${1:-$DB_NAME}" -N -s -e "$2" 2>&1
}

run_redis() {
    local auth_args=()
    [[ -n "$REDIS_AUTH" ]] && auth_args=(-a "$REDIS_AUTH" --no-auth-warning)
    redis-cli -h "$REDIS_HOST" -p "$REDIS_PORT" "${auth_args[@]}" "$@" 2>&1
}

# ===========================================================================
#  PostgreSQL tests
# ===========================================================================
run_postgresql_tests() {
    echo "${BOLD}PostgreSQL smoke tests — ${DB_HOST}:${DB_PORT}${RESET}"
    echo ""

    # 1. TCP connectivity
    run_test "TCP connectivity to ${DB_HOST}:${DB_PORT}" \
        check_port "$DB_HOST" "$DB_PORT"

    # 2. Authentication
    run_test "Authentication as ${DB_USER}" \
        bash -c 'PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c "SELECT 1" >/dev/null'

    # 3. Version
    run_test "Server version" \
        bash -c '
            ver=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c "SHOW server_version" 2>&1)
            echo "PostgreSQL $ver"
        '

    # 4. Database accessible
    run_test "Database '${DB_NAME}' accessible" \
        bash -c 'PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c "SELECT current_database()" >/dev/null'

    # 5. Replication lag
    if [[ "$SKIP_REPLICATION" == "true" ]]; then
        skip_test "Replication lag" "SKIP_REPLICATION=true"
    else
        run_test "Replication lag < ${MAX_REPLICATION_LAG_S}s" \
            bash -c '
                is_replica=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c "SELECT pg_is_in_recovery()" 2>&1)
                if [[ "$is_replica" == "t" ]]; then
                    lag=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c \
                        "SELECT COALESCE(EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))::int, 0)" 2>&1)
                    if [[ "$lag" -gt '"$MAX_REPLICATION_LAG_S"' ]]; then
                        echo "lag=${lag}s exceeds ${'"$MAX_REPLICATION_LAG_S"'}s"; exit 1
                    fi
                    echo "replica lag=${lag}s"
                else
                    echo "not a replica"; exit 2
                fi
            '
    fi

    # 6. Connection count
    run_test "Connection usage < ${MAX_CONNECTIONS_PCT}%" \
        bash -c '
            read -r used max_c <<< $(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c \
                "SELECT sum(numbackends), (SELECT setting::int FROM pg_settings WHERE name='"'"'max_connections'"'"') FROM pg_stat_database" 2>&1 | tr "|" " ")
            pct=$((used * 100 / max_c))
            if [[ $pct -ge '"$MAX_CONNECTIONS_PCT"' ]]; then
                echo "${used}/${max_c} (${pct}%)"; exit 1
            fi
            echo "${used}/${max_c} (${pct}%)"
        '

    # 7. Long-running queries
    run_test "No queries running > 300s" \
        bash -c '
            count=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c \
                "SELECT count(*) FROM pg_stat_activity WHERE state='"'"'active'"'"' AND now()-query_start > interval '"'"'300 seconds'"'"' AND pid <> pg_backend_pid()" 2>&1)
            if [[ "$count" -gt 0 ]]; then
                echo "${count} long-running queries found"; exit 1
            fi
            echo "none"
        '

    # 8. Table bloat
    run_test "Table bloat (dead tuple ratio < 20%)" \
        bash -c '
            worst=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c \
                "SELECT schemaname||'"'"'.'"'"'||relname||'"'"' '"'"'||round(100.0*n_dead_tup/(n_live_tup+n_dead_tup+1),1)||'"'"'%'"'"' FROM pg_stat_user_tables WHERE n_live_tup+n_dead_tup>1000 AND 100.0*n_dead_tup/(n_live_tup+n_dead_tup+1)>20 ORDER BY n_dead_tup DESC LIMIT 3" 2>&1)
            if [[ -n "$worst" ]]; then
                echo "bloated: $worst"; exit 1
            fi
            echo "ok"
        '

    # 9. Disk usage
    run_test "Disk usage for '${DB_NAME}'" \
        bash -c '
            size=$(PGPASSWORD="'"$DB_PASS"'" psql -h "'"$DB_HOST"'" -p "'"$DB_PORT"'" -U "'"$DB_USER"'" -d "'"$DB_NAME"'" -t -A -c \
                "SELECT pg_size_pretty(pg_database_size(current_database()))" 2>&1)
            echo "$size"
        '
}

# ===========================================================================
#  MySQL / MariaDB tests
# ===========================================================================
run_mysql_tests() {
    echo "${BOLD}MySQL smoke tests — ${DB_HOST}:${DB_PORT}${RESET}"
    echo ""

    # 1. TCP connectivity
    run_test "TCP connectivity to ${DB_HOST}:${DB_PORT}" \
        check_port "$DB_HOST" "$DB_PORT"

    # 2. Authentication
    run_test "Authentication as ${DB_USER}" \
        bash -c 'MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -e "SELECT 1" >/dev/null'

    # 3. Version
    run_test "Server version" \
        bash -c '
            ver=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SELECT version()" 2>&1)
            echo "MySQL $ver"
        '

    # 4. Database accessible
    run_test "Database '${DB_NAME}' accessible" \
        bash -c 'MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -D "'"$DB_NAME"'" -e "SELECT 1" >/dev/null'

    # 5. Replication lag
    if [[ "$SKIP_REPLICATION" == "true" ]]; then
        skip_test "Replication lag" "SKIP_REPLICATION=true"
    else
        run_test "Replication lag < ${MAX_REPLICATION_LAG_S}s" \
            bash -c '
                status=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SHOW REPLICA STATUS\G" 2>&1)
                if [[ -z "$status" ]]; then
                    status=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SHOW SLAVE STATUS\G" 2>&1)
                fi
                if [[ -z "$status" ]]; then
                    echo "not a replica"; exit 2
                fi
                lag=$(echo "$status" | grep -i "Seconds_Behind" | awk "{print \$NF}")
                if [[ "$lag" == "NULL" || -z "$lag" ]]; then
                    echo "replication not running (lag=NULL)"; exit 1
                fi
                if [[ "$lag" -gt '"$MAX_REPLICATION_LAG_S"' ]]; then
                    echo "lag=${lag}s exceeds '"$MAX_REPLICATION_LAG_S"'s"; exit 1
                fi
                echo "replica lag=${lag}s"
            '
    fi

    # 6. Connection count
    run_test "Connection usage < ${MAX_CONNECTIONS_PCT}%" \
        bash -c '
            used=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SELECT count(*) FROM information_schema.processlist" 2>&1)
            max_c=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SHOW VARIABLES LIKE '"'"'max_connections'"'"'" 2>&1 | awk "{print \$2}")
            pct=$((used * 100 / max_c))
            if [[ $pct -ge '"$MAX_CONNECTIONS_PCT"' ]]; then
                echo "${used}/${max_c} (${pct}%)"; exit 1
            fi
            echo "${used}/${max_c} (${pct}%)"
        '

    # 7. Slow query log
    run_test "Slow query log enabled" \
        bash -c '
            val=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SHOW VARIABLES LIKE '"'"'slow_query_log'"'"'" 2>&1 | awk "{print \$2}")
            if [[ "$val" != "ON" ]]; then
                echo "slow_query_log=$val"; exit 1
            fi
            echo "enabled"
        '

    # 8. Binary log space
    run_test "Binary log disk usage" \
        bash -c '
            logs=$(MYSQL_PWD="'"$DB_PASS"'" mysql -h "'"$DB_HOST"'" -P "'"$DB_PORT"'" -u "'"$DB_USER"'" -N -s -e "SHOW BINARY LOGS" 2>&1)
            if [[ "$logs" == *"not enabled"* || -z "$logs" ]]; then
                echo "binary logging disabled"; exit 2
            fi
            total=$(echo "$logs" | awk "{s+=\$2} END {printf \"%.1f MB\", s/1048576}")
            echo "$total"
        '
}

# ===========================================================================
#  Redis tests
# ===========================================================================
run_redis_tests() {
    echo "${BOLD}Redis smoke tests — ${REDIS_HOST}:${REDIS_PORT}${RESET}"
    echo ""

    local auth_args=()
    [[ -n "$REDIS_AUTH" ]] && auth_args=(-a "$REDIS_AUTH" --no-auth-warning)

    # 1. TCP connectivity
    run_test "TCP connectivity to ${REDIS_HOST}:${REDIS_PORT}" \
        check_port "$REDIS_HOST" "$REDIS_PORT"

    # 2. PING/PONG
    run_test "PING/PONG" \
        bash -c '
            reply=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' PING 2>&1)
            if [[ "$reply" != "PONG" ]]; then
                echo "got: $reply"; exit 1
            fi
            echo "PONG"
        '

    # 3. Server info
    run_test "Server info" \
        bash -c '
            info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO server 2>&1)
            ver=$(echo "$info" | grep "^redis_version:" | cut -d: -f2 | tr -d "\r")
            up=$(echo "$info" | grep "^uptime_in_days:" | cut -d: -f2 | tr -d "\r")
            echo "v${ver}, uptime ${up}d"
        '

    # 4. Memory usage
    run_test "Memory usage vs maxmemory" \
        bash -c '
            info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO memory 2>&1)
            used=$(echo "$info" | grep "^used_memory_human:" | cut -d: -f2 | tr -d "\r")
            max_raw=$(echo "$info" | grep "^maxmemory:" | cut -d: -f2 | tr -d "\r")
            max_h=$(echo "$info" | grep "^maxmemory_human:" | cut -d: -f2 | tr -d "\r")
            if [[ "$max_raw" == "0" ]]; then
                echo "used=${used}, maxmemory=unlimited"; exit 0
            fi
            echo "used=${used}, max=${max_h}"
        '

    # 5. Connected clients
    run_test "Connected clients" \
        bash -c '
            info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO clients 2>&1)
            count=$(echo "$info" | grep "^connected_clients:" | cut -d: -f2 | tr -d "\r")
            echo "${count} clients"
        '

    # 6. Replication status
    if [[ "$SKIP_REPLICATION" == "true" ]]; then
        skip_test "Replication status" "SKIP_REPLICATION=true"
    else
        run_test "Replication status" \
            bash -c '
                info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO replication 2>&1)
                role=$(echo "$info" | grep "^role:" | cut -d: -f2 | tr -d "\r")
                if [[ "$role" == "master" ]]; then
                    slaves=$(echo "$info" | grep "^connected_slaves:" | cut -d: -f2 | tr -d "\r")
                    echo "role=master, replicas=${slaves}"
                elif [[ "$role" == "slave" ]]; then
                    link=$(echo "$info" | grep "^master_link_status:" | cut -d: -f2 | tr -d "\r")
                    if [[ "$link" != "up" ]]; then
                        echo "replica link $link"; exit 1
                    fi
                    echo "role=replica, link=up"
                else
                    echo "role=$role"
                fi
            '
    fi

    # 7. Last save time
    if [[ "$SKIP_BACKUP_AGE" == "true" ]]; then
        skip_test "Last RDB/AOF save" "SKIP_BACKUP_AGE=true"
    else
        run_test "Last RDB save < ${MAX_BACKUP_AGE_H}h" \
            bash -c '
                info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO persistence 2>&1)
                last_save=$(echo "$info" | grep "^rdb_last_save_time:" | cut -d: -f2 | tr -d "\r")
                if [[ -z "$last_save" || "$last_save" == "0" ]]; then
                    echo "no RDB save recorded"; exit 2
                fi
                now=$(date +%s)
                age_h=$(( (now - last_save) / 3600 ))
                if [[ $age_h -gt '"$MAX_BACKUP_AGE_H"' ]]; then
                    echo "last save ${age_h}h ago (max '"$MAX_BACKUP_AGE_H"'h)"; exit 1
                fi
                echo "last save ${age_h}h ago"
            '
    fi

    # 8. Keyspace
    run_test "Keyspace info" \
        bash -c '
            info=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' INFO keyspace 2>&1)
            dbs=$(echo "$info" | grep "^db[0-9]" || true)
            if [[ -z "$dbs" ]]; then
                echo "no databases with keys"; exit 2
            fi
            total=0
            while IFS= read -r line; do
                keys=$(echo "$line" | grep -oP "keys=\K[0-9]+")
                total=$((total + keys))
            done <<< "$dbs"
            echo "${total} keys across $(echo "$dbs" | wc -l) database(s)"
        '

    # 9. Eviction policy
    run_test "Eviction policy" \
        bash -c '
            policy=$(redis-cli -h "'"$REDIS_HOST"'" -p "'"$REDIS_PORT"'" '"$(printf '%q ' "${auth_args[@]}")"' CONFIG GET maxmemory-policy 2>&1 | tail -1)
            echo "policy=$policy"
        '
}

# ===========================================================================
#  Summary
# ===========================================================================
print_summary() {
    local total=$((PASS_COUNT + FAIL_COUNT + SKIP_COUNT))
    echo ""
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "1..$total"
        echo "# pass  $PASS_COUNT"
        echo "# fail  $FAIL_COUNT"
        echo "# skip  $SKIP_COUNT"
    else
        echo "${BOLD}───────────────────────────────────────${RESET}"
        echo "  ${GREEN}PASS${RESET}  $PASS_COUNT"
        echo "  ${RED}FAIL${RESET}  $FAIL_COUNT"
        echo "  ${YELLOW}SKIP${RESET}  $SKIP_COUNT"
        echo "  Total $total"
        echo "${BOLD}───────────────────────────────────────${RESET}"
        if [[ $FAIL_COUNT -gt 0 ]]; then
            echo "  ${RED}${BOLD}RESULT: FAIL${RESET}"
        else
            echo "  ${GREEN}${BOLD}RESULT: PASS${RESET}"
        fi
    fi
}

# ===========================================================================
#  Main
# ===========================================================================
main() {
    if [[ -z "$DB_TYPE" ]]; then
        echo "ERROR: DB_TYPE is required (postgresql, mysql, redis)" >&2
        echo "Run with --help for usage information." >&2
        exit 1
    fi

    apply_defaults
    setup_colors

    case "$DB_TYPE" in
        postgresql) run_postgresql_tests ;;
        mysql)      run_mysql_tests ;;
        redis)      run_redis_tests ;;
        *)
            echo "ERROR: Unsupported DB_TYPE '${DB_TYPE}'" >&2
            exit 1
            ;;
    esac

    print_summary

    [[ $FAIL_COUNT -gt 0 ]] && exit 1
    exit 0
}

main "$@"
