#!/bin/bash
#
# Test script for vacuumdb --exclude-database feature
# Patch: "Add --exclude-database option to vacuumdb"
#
# Prerequisites:
#   - PostgreSQL built with the --exclude-database patch applied
#   - A running PostgreSQL instance (initdb + pg_ctl start)
#   - vacuumdb and psql in PATH (or specify with PGBIN)
#
# Usage:
#   ./test_exclude_database.sh                     # uses standard PG env vars
#   PGBIN=/path/to/bin ./test_exclude_database.sh  # custom binary path
#
# The script uses standard libpq environment variables (PGPORT, PGUSER,
# PGHOST, etc.) — same as psql, vacuumdb, and other PG tools.
# Set them before running if your instance uses non-default values.
#
# Optional:
#   PGBIN   - Path to PostgreSQL bin directory (default: uses PATH)
#   OUTFILE - Output file path (default: ./test_exclude_database_results.txt)
#
# The script will:
#   1. Create test databases (test_db1, test_db2, test_db3) if they don't exist
#   2. Run 11 tests covering the --exclude-database feature
#   3. Clean up (drops MixedCaseDB if created, keeps test_db1/2/3)
#   4. Write detailed results to OUTFILE
#
# ============================================================
# TEST PLAN — Description & Expected Results
# ============================================================
#
# TEST 1: Single database exclusion
#   Description: Run "vacuumdb --all --exclude-database=test_db1".
#                Verifies the basic functionality of excluding one
#                database from a vacuum --all operation.
#   Expected:    Exit code 0. test_db1 is NOT vacuumed.
#                All other connectable databases ARE vacuumed.
#
# TEST 2: Multiple database exclusions
#   Description: Run "vacuumdb --all -D test_db1 -D test_db2".
#                Verifies that the -D option can be specified multiple
#                times to exclude more than one database.
#   Expected:    Exit code 0. test_db1 and test_db2 are NOT vacuumed.
#                test_db3 and other databases ARE vacuumed.
#
# TEST 3: Exclude ALL databases (edge case)
#   Description: Run "vacuumdb --all -D <every connectable database>".
#                Tests the edge case where every database is excluded,
#                leaving nothing to vacuum.
#   Expected:    Exit code 0. No vacuum operations occur.
#                The command handles this gracefully without crashing.
#
# TEST 4: Exclude non-existent database
#   Description: Run "vacuumdb --all --exclude-database=nonexistent_db".
#                Tests that specifying a database name that doesn't exist
#                in pg_database is silently ignored.
#   Expected:    Exit code 0. All real databases are vacuumed normally.
#
# TEST 5: Exclude maintenance database (postgres)
#   Description: Run "vacuumdb --all --exclude-database=postgres".
#                Tests that the maintenance database itself can be excluded.
#                vacuumdb connects to postgres to retrieve the database list,
#                but should then skip it during the vacuum phase.
#   Expected:    Exit code 0. postgres is NOT vacuumed.
#                Other databases ARE vacuumed.
#
# TEST 6.1: Case sensitivity - exact case excludes
#   Description: Create "MixedCaseDB", then exclude with exact case.
#   Expected:    Exit code 0. MixedCaseDB is NOT vacuumed.
#
# TEST 6.2: Case sensitivity - wrong case does NOT exclude
#   Description: Exclude with wrong case "mixedcasedb".
#   Expected:    Exit code 0. MixedCaseDB IS vacuumed (case mismatch).
#
# TEST 7: --exclude-database without --all (error case)
#   Description: Run "vacuumdb -D test_db1" without --all.
#   Expected:    Non-zero exit code. Error message:
#                "cannot use --exclude-database without --all option"
#
# TEST 8: --exclude-database with -d (error case)
#   Description: Run "vacuumdb -d postgres -D test_db1".
#   Expected:    Non-zero exit code (conflicting options).
#
# TEST 9: SQL injection protection
#   Description: Pass a SQL injection attempt as the database name.
#   Expected:    Exit code 0. Safely escaped, all databases vacuumed.
#
# TEST 10: --help shows the new option
#   Description: Run "vacuumdb --help" and check the output.
#   Expected:    Output contains "-D, --exclude-database=DBNAME".
#
# TEST 11: -d with -D and --all (conflicting options)
#   Description: Run "vacuumdb -d postgres -D test_db1 --all".
#   Expected:    Non-zero exit code (-d and --all conflict).
#
# ============================================================

set -o pipefail

# ============================================================
# Configuration
# ============================================================
OUTFILE="${OUTFILE:-./test_exclude_database_results.txt}"

# Determine binary paths
if [ -n "$PGBIN" ]; then
    VACUUMDB="$PGBIN/vacuumdb"
    PSQL="$PGBIN/psql"
else
    VACUUMDB="$(command -v vacuumdb)"
    PSQL="$(command -v psql)"
fi

# Validate binaries exist
if [ ! -x "$VACUUMDB" ]; then
    echo "ERROR: vacuumdb not found. Set PGBIN or add to PATH."
    echo "  PGBIN=/path/to/pg/bin ./test_exclude_database.sh"
    exit 1
fi
if [ ! -x "$PSQL" ]; then
    echo "ERROR: psql not found. Set PGBIN or add to PATH."
    exit 1
fi

# Verify connection works
if ! $PSQL -d postgres -c "SELECT 1" >/dev/null 2>&1; then
    echo "ERROR: Cannot connect to PostgreSQL."
    echo "  Make sure PostgreSQL is running and standard PG env vars are set."
    echo "  Example: PGPORT=5433 PGUSER=myuser ./test_exclude_database.sh"
    exit 1
fi

# Verify --exclude-database option exists
if ! $VACUUMDB --help 2>&1 | grep -q "exclude-database"; then
    echo "ERROR: vacuumdb does not support --exclude-database."
    echo "  Make sure the patch is applied and vacuumdb is rebuilt."
    exit 1
fi

# ============================================================
# Test framework
# ============================================================
PASS_COUNT=0
FAIL_COUNT=0
TOTAL_COUNT=0
TEST_NUM=0

log() {
    echo "$1" | tee -a "$OUTFILE"
}

run_test() {
    local test_name="$1"
    local expected_result="$2"  # "pass" or "fail"
    local description="$3"
    shift 3
    local cmd="$@"

    TEST_NUM=$((TEST_NUM + 1))
    TOTAL_COUNT=$((TOTAL_COUNT + 1))
    log ""
    log "================================================================"
    log "TEST $TEST_NUM: $test_name"
    log "================================================================"
    log "Description: $description"
    log "Command: $cmd"
    log "Expected: $expected_result"
    log "---"

    LAST_OUTPUT=$(eval "$cmd" 2>&1)
    LAST_EXIT_CODE=$?

    log "Exit Code: $LAST_EXIT_CODE"
    log "Output:"
    log "$LAST_OUTPUT"
    log "---"

    if [ "$expected_result" = "pass" ]; then
        if [ $LAST_EXIT_CODE -eq 0 ]; then
            log "RESULT: ✅ PASS"
            PASS_COUNT=$((PASS_COUNT + 1))
            echo -e "${GREEN}✅ PASS${NC}: $test_name"
        else
            log "RESULT: ❌ FAIL (expected exit 0, got $LAST_EXIT_CODE)"
            FAIL_COUNT=$((FAIL_COUNT + 1))
            echo -e "${RED}❌ FAIL${NC}: $test_name"
        fi
    elif [ "$expected_result" = "fail" ]; then
        if [ $LAST_EXIT_CODE -ne 0 ]; then
            log "RESULT: ✅ PASS (correctly failed)"
            PASS_COUNT=$((PASS_COUNT + 1))
            echo -e "${GREEN}✅ PASS${NC}: $test_name"
        else
            log "RESULT: ❌ FAIL (expected non-zero exit, got 0)"
            FAIL_COUNT=$((FAIL_COUNT + 1))
            echo -e "${RED}❌ FAIL${NC}: $test_name"
        fi
    fi
}

check_output_contains() {
    local test_name="$1"
    local pattern="$2"
    local output="$3"

    TOTAL_COUNT=$((TOTAL_COUNT + 1))
    if echo "$output" | grep -q "$pattern"; then
        log "  ✅ '$pattern' found in output"
        PASS_COUNT=$((PASS_COUNT + 1))
        echo -e "${GREEN}  ✅${NC}: $test_name"
    else
        log "  ❌ '$pattern' NOT found in output"
        FAIL_COUNT=$((FAIL_COUNT + 1))
        echo -e "${RED}  ❌${NC}: $test_name"
    fi
}

check_output_not_contains() {
    local test_name="$1"
    local pattern="$2"
    local output="$3"

    TOTAL_COUNT=$((TOTAL_COUNT + 1))
    if echo "$output" | grep -q "$pattern"; then
        log "  ❌ '$pattern' found in output (should NOT be)"
        FAIL_COUNT=$((FAIL_COUNT + 1))
        echo -e "${RED}  ❌${NC}: $test_name"
    else
        log "  ✅ '$pattern' correctly absent from output"
        PASS_COUNT=$((PASS_COUNT + 1))
        echo -e "${GREEN}  ✅${NC}: $test_name"
    fi
}

# Colors (disabled if not a terminal)
if [ -t 1 ]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    NC='\033[0m'
else
    RED=''
    GREEN=''
    YELLOW=''
    NC=''
fi

# ============================================================
# SETUP
# ============================================================

> "$OUTFILE"

log "============================================================"
log "VACUUMDB --exclude-database FEATURE TEST SUITE"
log "============================================================"
log "Date: $(date)"
log "PostgreSQL: $($PSQL -d postgres -t -c 'SELECT version();' | head -1 | xargs)"
log "vacuumdb: $($VACUUMDB --version)"
log ""

# Create test databases if they don't exist
for db in test_db1 test_db2 test_db3; do
    $PSQL -d postgres -t -c \
        "SELECT 1 FROM pg_database WHERE datname='$db'" | grep -q 1 || \
        $PSQL -d postgres -c "CREATE DATABASE $db;" 2>/dev/null
done

log "Available databases:"
$PSQL -d postgres -t -c \
    "SELECT datname FROM pg_database WHERE datallowconn AND datconnlimit <> -2 ORDER BY 1;" | tee -a "$OUTFILE"
log ""

# ============================================================
# TEST 1: Single database exclusion
# ============================================================
run_test "Single database exclusion" "pass" \
    "Exclude test_db1 from vacuum --all. Expected: exit 0, test_db1 skipped, other dbs vacuumed." \
    "$VACUUMDB --all --exclude-database=test_db1"

check_output_not_contains "1.1: test_db1 NOT vacuumed" 'vacuuming database "test_db1"' "$LAST_OUTPUT"

# Verify all other databases were vacuumed
TOTAL_COUNT=$((TOTAL_COUNT + 1))
VACUUMED_COUNT=$(echo "$LAST_OUTPUT" | grep -c 'vacuuming database' || true)
TOTAL_DBS=$($PSQL -d postgres -t -c \
    "SELECT count(*) FROM pg_database WHERE datallowconn AND datconnlimit <> -2;" | tr -d ' ')
EXPECTED_COUNT=$((TOTAL_DBS - 1))
if [ "$VACUUMED_COUNT" -eq "$EXPECTED_COUNT" ]; then
    log "  ✅ All other databases vacuumed ($VACUUMED_COUNT/$EXPECTED_COUNT)"
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "${GREEN}  ✅${NC}: 1.2: All other databases vacuumed ($VACUUMED_COUNT/$EXPECTED_COUNT)"
else
    log "  ❌ Expected $EXPECTED_COUNT databases vacuumed, got $VACUUMED_COUNT"
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "${RED}  ❌${NC}: 1.2: Expected $EXPECTED_COUNT databases vacuumed, got $VACUUMED_COUNT"
fi

# ============================================================
# TEST 2: Multiple database exclusions
# ============================================================
run_test "Multiple database exclusions (-D db1 -D db2)" "pass" \
    "Exclude test_db1 and test_db2. Expected: exit 0, both excluded, test_db3 still vacuumed." \
    "$VACUUMDB --all -D test_db1 -D test_db2"

check_output_not_contains "2.1: test_db1 excluded" 'vacuuming database "test_db1"' "$LAST_OUTPUT"
check_output_not_contains "2.2: test_db2 excluded" 'vacuuming database "test_db2"' "$LAST_OUTPUT"
check_output_contains "2.3: test_db3 still vacuumed" 'vacuuming database "test_db3"' "$LAST_OUTPUT"

# ============================================================
# TEST 3: Exclude all databases (edge case)
# ============================================================
ALL_DBS=$($PSQL -d postgres -t -c \
    "SELECT datname FROM pg_database WHERE datallowconn AND datconnlimit <> -2;" | tr -d ' ' | grep -v '^$')
EXCLUDE_ARGS=""
for db in $ALL_DBS; do
    EXCLUDE_ARGS="$EXCLUDE_ARGS -D $db"
done

run_test "Exclude ALL databases (edge case)" "pass" \
    "Exclude every connectable database. Expected: exit 0, no vacuum ops, graceful no-op." \
    "$VACUUMDB --all $EXCLUDE_ARGS"

# ============================================================
# TEST 4: Exclude non-existent database
# ============================================================
run_test "Exclude non-existent database" "pass" \
    "Exclude a name that doesn't exist. Expected: exit 0, silently ignored, all real dbs vacuumed." \
    "$VACUUMDB --all --exclude-database=this_db_does_not_exist_xyz"

check_output_contains "4.1: postgres still vacuumed" 'vacuuming database "postgres"' "$LAST_OUTPUT"

# ============================================================
# TEST 5: Exclude maintenance database (postgres)
# ============================================================
run_test "Exclude maintenance database (postgres)" "pass" \
    "Exclude postgres (used for listing). Expected: exit 0, postgres skipped, others vacuumed." \
    "$VACUUMDB --all --exclude-database=postgres"

check_output_not_contains "5.1: postgres NOT vacuumed" 'vacuuming database "postgres"' "$LAST_OUTPUT"

# Verify all other databases were vacuumed
TOTAL_COUNT=$((TOTAL_COUNT + 1))
VACUUMED_COUNT=$(echo "$LAST_OUTPUT" | grep -c 'vacuuming database' || true)
EXPECTED_COUNT=$((TOTAL_DBS - 1))
if [ "$VACUUMED_COUNT" -eq "$EXPECTED_COUNT" ]; then
    log "  ✅ All other databases vacuumed ($VACUUMED_COUNT/$EXPECTED_COUNT)"
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "${GREEN}  ✅${NC}: 5.2: All other databases vacuumed ($VACUUMED_COUNT/$EXPECTED_COUNT)"
else
    log "  ❌ Expected $EXPECTED_COUNT databases vacuumed, got $VACUUMED_COUNT"
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "${RED}  ❌${NC}: 5.2: Expected $EXPECTED_COUNT databases vacuumed, got $VACUUMED_COUNT"
fi

# ============================================================
# TEST 6: Case sensitivity
# ============================================================
$PSQL -d postgres -c 'CREATE DATABASE "MixedCaseDB";' 2>/dev/null

# --- TEST 6.1: Exact case should exclude ---
run_test "Case sensitivity - exact case excludes MixedCaseDB" "pass" \
    "TEST 6.1: Exclude MixedCaseDB with exact case. Expected: exit 0, MixedCaseDB NOT vacuumed." \
    "$VACUUMDB --all --exclude-database=MixedCaseDB"

check_output_not_contains "6.1: MixedCaseDB excluded (exact case)" 'vacuuming database "MixedCaseDB"' "$LAST_OUTPUT"

# --- TEST 6.2: Wrong case should NOT exclude ---
log ""
log "----------------------------------------------------------------"
log "TEST 6.2: Case sensitivity - wrong case does NOT exclude"
log "----------------------------------------------------------------"
log "Description: Exclude with wrong case 'mixedcasedb'. Expected: MixedCaseDB IS vacuumed."
log "Command: $VACUUMDB --all --exclude-database=mixedcasedb"
log "---"

OUTPUT=$($VACUUMDB --all --exclude-database=mixedcasedb 2>&1)
log "Exit Code: $?"
log "Output:"
log "$OUTPUT"
log "---"

TOTAL_COUNT=$((TOTAL_COUNT + 1))
if echo "$OUTPUT" | grep -q 'vacuuming database "MixedCaseDB"'; then
    log "RESULT: ✅ PASS (MixedCaseDB was vacuumed — wrong case did not exclude it)"
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "${GREEN}✅ PASS${NC}: 6.2: MixedCaseDB IS vacuumed (wrong case doesn't match)"
else
    log "RESULT: ❌ FAIL (MixedCaseDB was excluded by wrong case — should not happen)"
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "${RED}❌ FAIL${NC}: 6.2: MixedCaseDB was incorrectly excluded by wrong case"
fi

# ============================================================
# TEST 7: --exclude-database without --all (error case)
# ============================================================
run_test "--exclude-database without --all (must fail)" "fail" \
    "Use -D without --all. Expected: non-zero exit, error: cannot use --exclude-database without --all." \
    "$VACUUMDB -D test_db1"

# Verify error message
TOTAL_COUNT=$((TOTAL_COUNT + 1))
if echo "$LAST_OUTPUT" | grep -q "cannot use --exclude-database without --all"; then
    log "  ✅ Correct error message: 'cannot use --exclude-database without --all'"
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "${GREEN}  ✅${NC}: 7.1: Correct error message"
else
    log "  ❌ Error message missing or wrong"
    log "  Got: $LAST_OUTPUT"
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "${RED}  ❌${NC}: 7.1: Error message missing or wrong"
fi

# ============================================================
# TEST 8: --exclude-database with -d (error case)
# ============================================================
run_test "--exclude-database with -d (must fail)" "fail" \
    "Use -D with -d. Expected: non-zero exit, -D requires --all but -d conflicts with --all." \
    "$VACUUMDB -d postgres -D test_db1"

# ============================================================
# TEST 9: SQL injection protection
# ============================================================
run_test "SQL injection attempt in database name" "pass" \
    "Inject SQL via -D value. Expected: exit 0, safely escaped, no injection, all dbs vacuumed." \
    "$VACUUMDB --all --exclude-database=\"'; DROP TABLE pg_class; --\""

# ============================================================
# TEST 10: --help shows the new option
# ============================================================
OUTPUT=$($VACUUMDB --help 2>&1)
TEST_NUM=$((TEST_NUM + 1))
TOTAL_COUNT=$((TOTAL_COUNT + 1))
log ""
log "================================================================"
log "TEST $TEST_NUM: --help shows --exclude-database"
log "================================================================"
log "Description: Run vacuumdb --help and check output."
log "Expected: Output contains '-D, --exclude-database=DBNAME'"
log "---"
log "Output:"
log "$OUTPUT"
log "---"
if echo "$OUTPUT" | grep -q "\-D.*--exclude-database"; then
    log "RESULT: ✅ PASS"
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "${GREEN}✅ PASS${NC}: --help shows --exclude-database"
else
    log "RESULT: ❌ FAIL"
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "${RED}❌ FAIL${NC}: --help shows --exclude-database"
fi

# ============================================================
# TEST 11: -d with -D and --all (conflicting options)
# ============================================================
run_test "-d with -D and --all (must fail)" "fail" \
    "Use -d postgres -D test_db1 --all. Expected: non-zero exit, -d and --all conflict." \
    "$VACUUMDB -d postgres -D test_db1 --all"

# ============================================================
# CLEANUP
# ============================================================
$PSQL -d postgres -c 'DROP DATABASE IF EXISTS "MixedCaseDB";' 2>/dev/null

# ============================================================
# SUMMARY
# ============================================================
log ""
log "============================================================"
log "TEST SUMMARY"
log "============================================================"
log "Total checks: $TOTAL_COUNT"
log "Passed:       $PASS_COUNT"
log "Failed:       $FAIL_COUNT"
log ""

if [ $FAIL_COUNT -eq 0 ]; then
    log "ALL TESTS PASSED"
    echo -e "\n${GREEN}ALL TESTS PASSED ($PASS_COUNT/$TOTAL_COUNT)${NC}"
else
    log "SOME TESTS FAILED"
    echo -e "\n${YELLOW}$FAIL_COUNT TESTS FAILED ($PASS_COUNT/$TOTAL_COUNT passed)${NC}"
fi

log ""
log "Results saved to: $OUTFILE"
echo ""
echo "Full results saved to: $OUTFILE"