#!/usr/bin/env bash
#
# Test suite for pg_scribe --start command
#
# This test suite:
# - Creates temporary test databases
# - Tests various --start scenarios
# - Verifies SQL capture (DML + DDL)
# - Tests signal handling
# - Cleans up all resources
#

set -euo pipefail

# Colors for test output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Test configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PG_SCRIBE="$SCRIPT_DIR/scripts/pg_scribe"
TEST_DIR="/tmp/pg_scribe_test_start_$$"
TEST_DB_PREFIX="pg_scribe_test_start_$$"
PGUSER="${PGUSER:-postgres}"

# Test counters
TESTS_RUN=0
TESTS_PASSED=0
TESTS_FAILED=0

# Cleanup tracking
DATABASES_TO_CLEANUP=()
PIDS_TO_CLEANUP=()

#
# Logging functions
#

log_test() {
    echo -e "${BLUE}TEST:${NC} $*"
}

log_pass() {
    echo -e "${GREEN}PASS:${NC} $*"
    ((TESTS_PASSED++))
}

log_fail() {
    echo -e "${RED}FAIL:${NC} $*"
    ((TESTS_FAILED++))
}

log_info() {
    echo -e "${YELLOW}INFO:${NC} $*"
}

#
# Helper functions
#

run_psql() {
    local dbname="$1"
    shift
    psql -U "$PGUSER" -d "$dbname" -tAq "$@"
}

query_db() {
    local dbname="$1"
    local query="$2"
    run_psql "$dbname" -c "$query" 2>/dev/null || true
}

create_test_db() {
    local dbname="$1"
    log_info "Creating test database: $dbname"

    # Drop if exists
    psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true

    # Create database
    psql -U "$PGUSER" -d postgres -c "CREATE DATABASE $dbname;" &>/dev/null

    DATABASES_TO_CLEANUP+=("$dbname")
}

# shellcheck disable=SC2317  # Function called from cleanup trap handler
drop_test_db() {
    local dbname="$1"
    log_info "Dropping test database: $dbname"

    # Terminate connections
    psql -U "$PGUSER" -d postgres -c "
        SELECT pg_terminate_backend(pid)
        FROM pg_stat_activity
        WHERE datname = '$dbname' AND pid <> pg_backend_pid();
    " &>/dev/null || true

    # Drop database
    psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true
}

# shellcheck disable=SC2317  # Function called from cleanup trap handler
drop_replication_slot() {
    local dbname="$1"
    local slot="$2"
    log_info "Dropping replication slot: $slot"

    # Check if slot exists
    local exists
    exists=$(query_db "$dbname" "
        SELECT 1 FROM pg_replication_slots WHERE slot_name = '$slot';
    ")

    if [[ -n "$exists" ]]; then
        # Drop slot
        query_db "$dbname" "SELECT pg_drop_replication_slot('$slot');" || true
    fi
}

create_table_with_pk() {
    local dbname="$1"
    local table="$2"
    query_db "$dbname" "
        CREATE TABLE $table (
            id SERIAL PRIMARY KEY,
            name TEXT,
            created_at TIMESTAMP DEFAULT now()
        );
    "
}

initialize_backup_system() {
    local dbname="$1"
    local slot="$2"
    local backup_dir="$3"

    # Create backup directory
    mkdir -p "$backup_dir"

    # Initialize
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null
}

#
# Test cases
#

test_start_without_init() {
    ((TESTS_RUN++))
    log_test "Start without initialization (should fail)"

    local dbname="${TEST_DB_PREFIX}_noinit"
    local slot="test_slot_noinit"
    local backup_dir="$TEST_DIR/noinit"

    # Setup - create db and backup dir but DON'T initialize
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    mkdir -p "$backup_dir"

    # Try to start - should fail with exit code 4 (no chains found)
    local exit_code=0
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 4 ]]; then
        log_pass "Correctly failed with backup error (no chains)"
        return 0
    else
        log_fail "Expected exit code 4, got $exit_code"
        return 1
    fi
}

test_start_basic_streaming() {
    ((TESTS_RUN++))
    log_test "Basic streaming with DML capture"

    local dbname="${TEST_DB_PREFIX}_basic"
    local slot="test_slot_basic"
    local backup_dir="$TEST_DIR/basic"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Make some changes
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Alice');"
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Bob');"
    query_db "$dbname" "UPDATE users SET name = 'Alice Smith' WHERE name = 'Alice';"
    query_db "$dbname" "DELETE FROM users WHERE name = 'Bob';"

    # Give it time to flush (status-interval=1, fsync-interval=1)
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Find active.sql in chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local output_file="${chain_dirs[0]}/active.sql"

    # Verify output file exists
    if [[ ! -f "$output_file" ]]; then
        log_fail "Output file not created: $output_file"
        return 1
    fi

    # Verify SQL content
    if ! grep -q "INSERT INTO public.users" "$output_file"; then
        log_fail "INSERT not captured"
        return 1
    fi

    if ! grep -q "UPDATE public.users" "$output_file"; then
        log_fail "UPDATE not captured"
        return 1
    fi

    if ! grep -q "DELETE FROM public.users" "$output_file"; then
        log_fail "DELETE not captured"
        return 1
    fi

    # Verify transaction boundaries
    if ! grep -q "BEGIN" "$output_file"; then
        log_fail "BEGIN not captured"
        return 1
    fi

    if ! grep -q "COMMIT" "$output_file"; then
        log_fail "COMMIT not captured"
        return 1
    fi

    log_pass "DML captured successfully"
    return 0
}

test_start_ddl_capture() {
    ((TESTS_RUN++))
    log_test "DDL capture via event triggers"

    local dbname="${TEST_DB_PREFIX}_ddl"
    local slot="test_slot_ddl"
    local backup_dir="$TEST_DIR/ddl"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Make DDL changes
    query_db "$dbname" "CREATE TABLE products (id SERIAL PRIMARY KEY, name TEXT);"
    query_db "$dbname" "ALTER TABLE products ADD COLUMN price NUMERIC(10,2);"
    query_db "$dbname" "DROP TABLE products;"

    # Give it time to flush
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Find active.sql in chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local output_file="${chain_dirs[0]}/active.sql"

    # Verify DDL captured
    if ! grep -qi "CREATE TABLE products" "$output_file"; then
        log_fail "CREATE TABLE not captured"
        return 1
    fi

    if ! grep -qi "ALTER TABLE products" "$output_file"; then
        log_fail "ALTER TABLE not captured"
        return 1
    fi

    if ! grep -qi "DROP TABLE products" "$output_file"; then
        log_fail "DROP TABLE not captured"
        return 1
    fi

    log_pass "DDL captured successfully"
    return 0
}

test_start_truncate_capture() {
    ((TESTS_RUN++))
    log_test "TRUNCATE capture"

    local dbname="${TEST_DB_PREFIX}_truncate"
    local slot="test_slot_truncate"
    local backup_dir="$TEST_DIR/truncate"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Insert data and truncate
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Alice'), ('Bob'), ('Charlie');"
    query_db "$dbname" "TRUNCATE users;"

    # Give it time to flush
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Find active.sql in chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local output_file="${chain_dirs[0]}/active.sql"

    # Verify TRUNCATE captured
    if ! grep -qi "TRUNCATE.*users" "$output_file"; then
        log_fail "TRUNCATE not captured"
        return 1
    fi

    log_pass "TRUNCATE captured successfully"
    return 0
}

test_start_signal_handling() {
    ((TESTS_RUN++))
    log_test "Signal handling (SIGTERM)"

    local dbname="${TEST_DB_PREFIX}_signal"
    local slot="test_slot_signal"
    local backup_dir="$TEST_DIR/signal"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Make a change
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Test');"

    # Give it time to flush
    sleep 1

    # Send SIGTERM
    kill -TERM "$pg_scribe_pid" 2>/dev/null || true

    # Wait for graceful shutdown (with timeout)
    local timeout=3
    local count=0
    while kill -0 "$pg_scribe_pid" 2>/dev/null && [[ $count -lt $timeout ]]; do
        sleep 1
        ((count++))
    done

    # Check if process stopped
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        log_fail "Process did not stop after SIGTERM"
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
        return 1
    fi

    # Find active.sql in chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local output_file="${chain_dirs[0]}/active.sql"

    # Verify output file was created and flushed
    if [[ ! -f "$output_file" ]]; then
        log_fail "Output file not created"
        return 1
    fi

    if ! grep -q "INSERT INTO public.users" "$output_file"; then
        log_fail "Data not flushed before shutdown"
        return 1
    fi

    log_pass "SIGTERM handled gracefully"
    return 0
}

test_start_interleaved_ddl_dml() {
    ((TESTS_RUN++))
    log_test "DDL and DML interleaving (chronological order)"

    local dbname="${TEST_DB_PREFIX}_interleaved"
    local slot="test_slot_interleaved"
    local backup_dir="$TEST_DIR/interleaved"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Make interleaved changes
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Alice');"
    query_db "$dbname" "ALTER TABLE users ADD COLUMN email TEXT;"
    query_db "$dbname" "UPDATE users SET email = 'alice@example.com' WHERE name = 'Alice';"
    query_db "$dbname" "ALTER TABLE users DROP COLUMN created_at;"
    query_db "$dbname" "INSERT INTO users (name, email) VALUES ('Bob', 'bob@example.com');"

    # Give it time to flush
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Find active.sql in chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local output_file="${chain_dirs[0]}/active.sql"

    # Verify all operations captured
    if ! grep -q "INSERT INTO public.users (id, name" "$output_file"; then
        log_fail "First INSERT not captured"
        return 1
    fi

    if ! grep -qi "ALTER TABLE.*ADD COLUMN email" "$output_file"; then
        log_fail "ALTER TABLE ADD COLUMN not captured"
        return 1
    fi

    if ! grep -q "UPDATE public.users.*email" "$output_file"; then
        log_fail "UPDATE with new column not captured"
        return 1
    fi

    if ! grep -qi "ALTER TABLE.*DROP COLUMN" "$output_file"; then
        log_fail "ALTER TABLE DROP COLUMN not captured"
        return 1
    fi

    # Verify chronological order by checking line numbers
    local insert1_line
    insert1_line=$(grep -n "INSERT INTO public.users (id, name" "$output_file" | head -1 | cut -d: -f1)
    local alter_add_line
    alter_add_line=$(grep -ni "ALTER TABLE.*ADD COLUMN email" "$output_file" | cut -d: -f1)
    local update_line
    update_line=$(grep -n "UPDATE public.users.*email" "$output_file" | cut -d: -f1)

    if [[ "$insert1_line" -gt "$alter_add_line" ]]; then
        log_fail "DDL/DML not in chronological order (INSERT after ALTER)"
        return 1
    fi

    if [[ "$alter_add_line" -gt "$update_line" ]]; then
        log_fail "DDL/DML not in chronological order (ALTER after UPDATE)"
        return 1
    fi

    log_pass "DDL/DML interleaved in correct chronological order"
    return 0
}

test_rotate_diff_basic() {
    ((TESTS_RUN++))
    log_test "--rotate-diff basic functionality"

    local dbname="${TEST_DB_PREFIX}_rotate_basic"
    local slot="test_slot_rotate_basic"
    local backup_dir="$TEST_DIR/rotate_basic"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    # Make changes before rotation
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Before Rotation');"

    # Give it time to flush
    sleep 2

    # Find chain directory
    local chain_dirs=("$backup_dir"/chain-*)
    local chain_dir="${chain_dirs[0]}"

    # Verify active.sql exists
    if [[ ! -f "$chain_dir/active.sql" ]]; then
        log_fail "active.sql not found before rotation"
        return 1
    fi

    # Perform rotation
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null

    # Give pg_recvlogical time to reopen the file
    sleep 2

    # Verify sealed differential was created
    local diff_files=("$chain_dir"/diff-*.sql)
    if [[ ! -f "${diff_files[0]}" ]]; then
        log_fail "Sealed differential not created"
        return 1
    fi

    # Verify new active.sql was created
    if [[ ! -f "$chain_dir/active.sql" ]]; then
        log_fail "New active.sql not created after rotation"
        return 1
    fi

    # Make changes after rotation
    query_db "$dbname" "INSERT INTO users (name) VALUES ('After Rotation');"

    # Give it time to flush
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Verify sealed differential has "Before Rotation" but not "After Rotation"
    if ! grep -q "Before Rotation" "${diff_files[0]}"; then
        log_fail "Sealed differential missing data before rotation"
        return 1
    fi

    if grep -q "After Rotation" "${diff_files[0]}"; then
        log_fail "Sealed differential should not contain data after rotation"
        return 1
    fi

    # Verify new active.sql has "After Rotation" but not "Before Rotation"
    if ! grep -q "After Rotation" "$chain_dir/active.sql"; then
        log_fail "New active.sql missing data after rotation"
        return 1
    fi

    if grep -q "Before Rotation" "$chain_dir/active.sql"; then
        log_fail "New active.sql should not contain data before rotation"
        return 1
    fi

    log_pass "--rotate-diff works correctly"
    return 0
}

test_rotate_diff_no_active_process() {
    ((TESTS_RUN++))
    log_test "--rotate-diff without active process (should fail)"

    local dbname="${TEST_DB_PREFIX}_rotate_noactive"
    local slot="test_slot_rotate_noactive"
    local backup_dir="$TEST_DIR/rotate_noactive"

    # Setup - initialize but don't start streaming
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Try to rotate without active process - should fail
    local exit_code=0
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null || exit_code=$?

    if [[ $exit_code -ne 0 ]]; then
        log_pass "Correctly failed when no active process"
        return 0
    else
        log_fail "Should have failed with no active process"
        return 1
    fi
}

test_rotate_diff_stale_pidfile() {
    ((TESTS_RUN++))
    log_test "--rotate-diff with stale pidfile (should fail)"

    local dbname="${TEST_DB_PREFIX}_rotate_stale"
    local slot="test_slot_rotate_stale"
    local backup_dir="$TEST_DIR/rotate_stale"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Create a stale pidfile (process doesn't exist)
    echo "99999" > "$backup_dir/.pg_scribe.pid"

    # Try to rotate - should fail with stale pidfile
    local exit_code=0
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null || exit_code=$?

    if [[ $exit_code -ne 0 ]]; then
        log_pass "Correctly failed with stale pidfile"
        return 0
    else
        log_fail "Should have failed with stale pidfile"
        return 1
    fi
}

test_rotate_diff_multiple() {
    ((TESTS_RUN++))
    log_test "--rotate-diff multiple times"

    local dbname="${TEST_DB_PREFIX}_rotate_multiple"
    local slot="test_slot_rotate_multiple"
    local backup_dir="$TEST_DIR/rotate_multiple"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Start streaming in background
    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" -s 1 -F 1 &>/dev/null &
    local pg_scribe_pid=$!
    PIDS_TO_CLEANUP+=("$pg_scribe_pid")

    # Give it a moment to start
    sleep 1

    local chain_dirs=("$backup_dir"/chain-*)
    local chain_dir="${chain_dirs[0]}"

    # First rotation
    query_db "$dbname" "INSERT INTO users (name) VALUES ('First');"
    sleep 2
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null
    sleep 2

    # Second rotation
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Second');"
    sleep 2
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null
    sleep 2

    # Third rotation
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Third');"
    sleep 2
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" &>/dev/null
    sleep 2

    # Stop streaming
    kill -INT "$pg_scribe_pid" 2>/dev/null || true
    sleep 1

    # Force kill if still running
    if kill -0 "$pg_scribe_pid" 2>/dev/null; then
        kill -9 "$pg_scribe_pid" 2>/dev/null || true
    fi
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Verify we have 3 sealed differentials
    local diff_count
    diff_count=$(find "$chain_dir" -maxdepth 1 -name 'diff-*.sql' -type f 2>/dev/null | wc -l)

    if [[ $diff_count -ne 3 ]]; then
        log_fail "Expected 3 differentials, found $diff_count"
        return 1
    fi

    # Verify each differential has the correct data
    local diff_files=("$chain_dir"/diff-*.sql)

    if ! grep -q "First" "${diff_files[0]}"; then
        log_fail "First differential missing expected data"
        return 1
    fi

    if ! grep -q "Second" "${diff_files[1]}"; then
        log_fail "Second differential missing expected data"
        return 1
    fi

    if ! grep -q "Third" "${diff_files[2]}"; then
        log_fail "Third differential missing expected data"
        return 1
    fi

    # Verify active.sql exists (from streaming after last rotation)
    if [[ ! -f "$chain_dir/active.sql" ]]; then
        log_fail "active.sql not found after multiple rotations"
        return 1
    fi

    log_pass "Multiple rotations work correctly"
    return 0
}

#
# Cleanup
#

# shellcheck disable=SC2317  # Function called via trap handler
cleanup() {
    log_info "Cleaning up test resources..."

    # Kill any running pg_scribe processes
    for pid in "${PIDS_TO_CLEANUP[@]}"; do
        if kill -0 "$pid" 2>/dev/null; then
            log_info "Stopping pg_scribe process $pid"
            # Try graceful shutdown first (allows signal forwarding to child processes)
            kill -TERM "$pid" 2>/dev/null || true

            # Wait briefly for graceful shutdown
            local timeout=2
            local count=0
            while kill -0 "$pid" 2>/dev/null && [[ $count -lt $timeout ]]; do
                sleep 0.5
                ((count++))
            done

            # Force kill if still running
            if kill -0 "$pid" 2>/dev/null; then
                log_info "Force killing pg_scribe process $pid"
                kill -9 "$pid" 2>/dev/null || true
            fi
        fi
    done

    # Wait for child pg_recvlogical processes to fully terminate
    # (They may take a moment to shut down after parent terminates)
    sleep 1

    # Drop replication slots
    for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
        for slot in test_slot_noinit test_slot_basic test_slot_ddl test_slot_truncate test_slot_signal test_slot_interleaved test_slot_rotate_basic test_slot_rotate_noactive test_slot_rotate_stale test_slot_rotate_multiple; do
            drop_replication_slot "$dbname" "$slot" 2>/dev/null || true
        done
    done

    # Drop databases
    for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
        drop_test_db "$dbname"
    done

    # Remove test directory
    if [[ -d "$TEST_DIR" ]]; then
        rm -rf "$TEST_DIR"
    fi

    log_info "Cleanup complete"
}

#
# Main test runner
#

main() {
    echo "========================================"
    echo "pg_scribe --start Test Suite"
    echo "========================================"
    echo ""

    # Verify pg_scribe exists
    if [[ ! -x "$PG_SCRIBE" ]]; then
        echo "ERROR: pg_scribe not found or not executable: $PG_SCRIBE"
        exit 1
    fi

    # Verify PostgreSQL is running
    if ! psql -U "$PGUSER" -d postgres -c "SELECT 1;" &>/dev/null; then
        echo "ERROR: Cannot connect to PostgreSQL"
        exit 1
    fi

    # Verify wal_level is logical
    local wal_level
    wal_level=$(psql -U "$PGUSER" -d postgres -tAq -c "SHOW wal_level;")
    if [[ "$wal_level" != "logical" ]]; then
        echo "ERROR: wal_level must be 'logical', currently: $wal_level"
        echo "Update ~/.pgenv/pgsql/data/postgresql.conf and restart PostgreSQL"
        exit 1
    fi

    # Verify wal2sql extension is available
    if ! pg_config --pkglibdir &>/dev/null; then
        echo "ERROR: pg_config not found"
        exit 1
    fi

    local wal2sql_path
    wal2sql_path="$(pg_config --pkglibdir)/wal2sql.so"
    if [[ ! -f "$wal2sql_path" ]]; then
        echo "ERROR: wal2sql.so not found at $wal2sql_path"
        echo "Build and install wal2sql: cd wal2sql && make && make install"
        exit 1
    fi

    # Create test directory
    mkdir -p "$TEST_DIR"

    # Set up cleanup trap
    trap cleanup EXIT INT TERM

    echo "Running tests..."
    echo ""

    # Run all tests (use || true to prevent set -e from exiting)
    test_start_without_init || true
    test_start_basic_streaming || true
    test_start_ddl_capture || true
    test_start_truncate_capture || true
    test_start_signal_handling || true
    test_start_interleaved_ddl_dml || true
    test_rotate_diff_basic || true
    test_rotate_diff_no_active_process || true
    test_rotate_diff_stale_pidfile || true
    test_rotate_diff_multiple || true

    # Summary
    echo ""
    echo "========================================"
    echo "Test Results"
    echo "========================================"
    echo "Tests run:    $TESTS_RUN"
    echo -e "Tests passed: ${GREEN}$TESTS_PASSED${NC}"
    echo -e "Tests failed: ${RED}$TESTS_FAILED${NC}"
    echo ""

    if [[ $TESTS_FAILED -eq 0 ]]; then
        echo -e "${GREEN}All tests passed!${NC}"
        exit 0
    else
        echo -e "${RED}Some tests failed!${NC}"
        exit 1
    fi
}

main "$@"