#!/usr/bin/env bash
#
# Test suite for pg_scribe --restore command
#
# This test suite:
# - Creates temporary test databases
# - Tests various --restore scenarios
# - Verifies expected outcomes
# - 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_restore_$$"
TEST_DB_PREFIX="pg_scribe_restore_test_$$"
PGUSER="${PGUSER:-postgres}"

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

# Cleanup tracking
DATABASES_TO_CLEANUP=()
SLOTS_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() {
    local dbname="$1"
    local backup_dir="$2"
    local slot="$3"

    # Init backup system
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null
    SLOTS_TO_CLEANUP+=("$dbname:$slot")
}

#
# Test cases
#

test_basic_restore_from_base_only() {
    ((TESTS_RUN++))
    log_test "Basic restore from base backup only"

    local source_db="${TEST_DB_PREFIX}_source_basic"
    local restore_db="${TEST_DB_PREFIX}_restore_basic"
    local slot="test_slot_basic"
    local backup_dir="$TEST_DIR/basic"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "users"
    query_db "$source_db" "INSERT INTO users (name) VALUES ('Alice'), ('Bob'), ('Charlie');"

    # Initialize backup
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Restore to new database with --create
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Verify data
        local count
        count=$(query_db "$restore_db" "SELECT COUNT(*) FROM users;")
        if [[ "$count" -ne 3 ]]; then
            log_fail "Expected 3 rows, got $count"
            return 1
        fi

        # Verify specific data (psql returns one value per line)
        local names
        names=$(query_db "$restore_db" "SELECT name FROM users ORDER BY id;")
        local expected_names=$'Alice\nBob\nCharlie'
        if [[ "$names" != "$expected_names" ]]; then
            log_fail "Data mismatch: expected '$expected_names', got '$names'"
            return 1
        fi

        log_pass "Basic restore successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_with_differentials() {
    ((TESTS_RUN++))
    log_test "Restore with base + sealed differentials"

    local source_db="${TEST_DB_PREFIX}_source_diff"
    local restore_db="${TEST_DB_PREFIX}_restore_diff"
    local slot="test_slot_diff"
    local backup_dir="$TEST_DIR/differentials"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "orders"
    query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order1'), ('Order2');"

    # Initialize backup (creates first chain)
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Start streaming to latest chain
    "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
    local pg_scribe_pid=$!

    # Wait for streaming to start
    sleep 2

    # Make some changes
    query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order3');"
    query_db "$source_db" "UPDATE orders SET name = 'Order1_Updated' WHERE name = 'Order1';"
    query_db "$source_db" "DELETE FROM orders WHERE name = 'Order2';"

    # Wait for changes to be captured
    sleep 2

    # Rotate differential to seal it
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
    sleep 1

    # Stop streaming
    kill -TERM "$pg_scribe_pid" 2>/dev/null || true
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Find latest chain
    local chain_dir
    chain_dir=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | sort | tail -1)

    # Verify sealed differential exists
    local diff_count
    diff_count=$(find "$chain_dir" -name 'diff-*.sql' 2>/dev/null | wc -l)
    if [[ "$diff_count" -eq 0 ]]; then
        log_fail "No sealed differentials found"
        return 1
    fi

    # Restore to new database
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Verify final state (should have Order1_Updated and Order3, not Order2)
        local count
        count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
        if [[ "$count" -ne 2 ]]; then
            log_fail "Expected 2 rows after differential restore, got $count"
            return 1
        fi

        # Verify Order1 was updated
        local order1_name
        order1_name=$(query_db "$restore_db" "SELECT name FROM orders WHERE id = 1;")
        if [[ "$order1_name" != "Order1_Updated" ]]; then
            log_fail "UPDATE not applied: expected 'Order1_Updated', got '$order1_name'"
            return 1
        fi

        # Verify Order3 exists
        local order3_exists
        order3_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order3';")
        if [[ "$order3_exists" -ne 1 ]]; then
            log_fail "INSERT not applied: Order3 not found"
            return 1
        fi

        # Verify Order2 was deleted
        local order2_exists
        order2_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order2';")
        if [[ "$order2_exists" -ne 0 ]]; then
            log_fail "DELETE not applied: Order2 still exists"
            return 1
        fi

        log_pass "Restore with differentials successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_with_ddl_changes() {
    ((TESTS_RUN++))
    log_test "Restore with DDL changes in differential"

    local source_db="${TEST_DB_PREFIX}_source_ddl"
    local restore_db="${TEST_DB_PREFIX}_restore_ddl"
    local slot="test_slot_ddl"
    local backup_dir="$TEST_DIR/ddl"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "products"
    query_db "$source_db" "INSERT INTO products (name) VALUES ('Widget');"

    # Initialize backup (creates chain)
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Start streaming
    "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
    local pg_scribe_pid=$!

    # Wait for streaming to start
    sleep 2

    # Add a column (DDL change)
    query_db "$source_db" "ALTER TABLE products ADD COLUMN price NUMERIC(10,2);"

    # Insert data using new column
    query_db "$source_db" "INSERT INTO products (name, price) VALUES ('Gadget', 19.99);"

    # Wait for changes to be captured
    sleep 2

    # Rotate to seal differential
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
    sleep 1

    # Stop streaming
    kill -TERM "$pg_scribe_pid" 2>/dev/null || true
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Restore to new database
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Verify new column exists
        local has_price_column
        has_price_column=$(query_db "$restore_db" "
            SELECT COUNT(*) FROM information_schema.columns
            WHERE table_name = 'products' AND column_name = 'price';
        ")
        if [[ "$has_price_column" -ne 1 ]]; then
            log_fail "DDL not applied: price column not found"
            return 1
        fi

        # Verify data with new column
        local gadget_price
        gadget_price=$(query_db "$restore_db" "SELECT price FROM products WHERE name = 'Gadget';")
        if [[ "$gadget_price" != "19.99" ]]; then
            log_fail "Data with new column not correct: expected '19.99', got '$gadget_price'"
            return 1
        fi

        log_pass "Restore with DDL changes successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_sequence_synchronization() {
    ((TESTS_RUN++))
    log_test "Restore with sequence synchronization"

    local source_db="${TEST_DB_PREFIX}_source_seq"
    local restore_db="${TEST_DB_PREFIX}_restore_seq"
    local slot="test_slot_seq"
    local backup_dir="$TEST_DIR/sequence"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "items"
    query_db "$source_db" "INSERT INTO items (name) VALUES ('Item1'), ('Item2'), ('Item3');"

    # Initialize backup (creates first chain)
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Add more data to advance sequence
    query_db "$source_db" "INSERT INTO items (name) VALUES ('Item4'), ('Item5');"

    # Start streaming to capture incremental changes
    "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
    local pg_scribe_pid=$!
    sleep 2

    # Rotate differential to seal it
    "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
    sleep 1

    # Stop streaming
    kill -TERM "$pg_scribe_pid" 2>/dev/null || true
    wait "$pg_scribe_pid" 2>/dev/null || true

    # Restore to new database
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Get current sequence value
        local seq_val
        seq_val=$(query_db "$restore_db" "SELECT last_value FROM items_id_seq;")

        # Should be at least 5 (we inserted 5 items total)
        if [[ "$seq_val" -lt 5 ]]; then
            log_fail "Sequence not synchronized: expected >= 5, got $seq_val"
            return 1
        fi

        # Try inserting new row - should get ID 6
        query_db "$restore_db" "INSERT INTO items (name) VALUES ('Item6');"
        local new_id
        new_id=$(query_db "$restore_db" "SELECT id FROM items WHERE name = 'Item6';")
        if [[ "$new_id" -ne 6 ]]; then
            log_fail "Next sequence value incorrect: expected 6, got $new_id"
            return 1
        fi

        log_pass "Sequence synchronization successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_no_sync_sequences() {
    ((TESTS_RUN++))
    log_test "Restore with --no-sync-sequences flag"

    local source_db="${TEST_DB_PREFIX}_source_noseq"
    local restore_db="${TEST_DB_PREFIX}_restore_noseq"
    local slot="test_slot_noseq"
    local backup_dir="$TEST_DIR/noseq"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "records"
    query_db "$source_db" "INSERT INTO records (name) VALUES ('Rec1'), ('Rec2');"

    # Initialize backup (creates chain)
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Restore with --no-sync-sequences
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" --no-sync-sequences &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Sequence should be at the value from pg_dump (2)
        local seq_val
        seq_val=$(query_db "$restore_db" "SELECT last_value FROM records_id_seq;")

        # With --no-sync-sequences, sequence value is from pg_dump
        # It should match what's in the backup
        if [[ "$seq_val" -lt 1 ]]; then
            log_fail "Sequence value unexpectedly low: $seq_val"
            return 1
        fi

        log_pass "Restore with --no-sync-sequences successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_to_existing_database() {
    ((TESTS_RUN++))
    log_test "Restore to existing database (without --create)"

    local source_db="${TEST_DB_PREFIX}_source_exist"
    local restore_db="${TEST_DB_PREFIX}_restore_exist"
    local slot="test_slot_exist"
    local backup_dir="$TEST_DIR/existing"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "data"
    query_db "$source_db" "INSERT INTO data (name) VALUES ('Test1');"

    # Initialize backup
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Pre-create target database
    create_test_db "$restore_db"

    # Restore without --create
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null; then
        # Verify data
        local count
        count=$(query_db "$restore_db" "SELECT COUNT(*) FROM data;")
        if [[ "$count" -ne 1 ]]; then
            log_fail "Expected 1 row, got $count"
            return 1
        fi

        log_pass "Restore to existing database successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_fails_if_db_exists_with_create() {
    ((TESTS_RUN++))
    log_test "Restore fails if database exists with --create"

    local source_db="${TEST_DB_PREFIX}_source_exists"
    local restore_db="${TEST_DB_PREFIX}_restore_exists"
    local slot="test_slot_exists"
    local backup_dir="$TEST_DIR/exists"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "test"
    query_db "$source_db" "INSERT INTO test (name) VALUES ('Test');"

    # Initialize backup
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Pre-create target database
    create_test_db "$restore_db"

    # Restore with --create should fail
    local exit_code=0
    "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 4 ]]; then
        log_pass "Correctly failed when database exists with --create"
        return 0
    else
        log_fail "Expected exit code 4, got $exit_code"
        return 1
    fi
}

test_restore_missing_backup_directory() {
    ((TESTS_RUN++))
    log_test "Restore fails with missing backup directory"

    local restore_db="${TEST_DB_PREFIX}_restore_missing"
    local backup_dir="$TEST_DIR/nonexistent"

    # Try to restore from non-existent directory
    local exit_code=0
    "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null || exit_code=$?

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

test_restore_specific_chain() {
    ((TESTS_RUN++))
    log_test "Restore specific chain with --base-backup (chain ID)"

    local source_db="${TEST_DB_PREFIX}_source_specific"
    local restore_db="${TEST_DB_PREFIX}_restore_specific"
    local slot="test_slot_specific"
    local backup_dir="$TEST_DIR/specific"

    # Setup source database
    create_test_db "$source_db"
    create_table_with_pk "$source_db" "entries"
    query_db "$source_db" "INSERT INTO entries (name) VALUES ('Entry1');"

    # Initialize backup (creates first chain)
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Find the chain ID
    local chain_dir
    chain_dir=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | sort | tail -1)

    if [[ -z "$chain_dir" ]]; then
        log_fail "Chain directory not found"
        return 1
    fi

    local chain_id
    chain_id=$(basename "$chain_dir" | sed 's/^chain-//')

    # Restore using specific chain ID
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" --base-backup="$chain_id" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Verify data
        local count
        count=$(query_db "$restore_db" "SELECT COUNT(*) FROM entries;")
        if [[ "$count" -ne 1 ]]; then
            log_fail "Expected 1 row, got $count"
            return 1
        fi

        log_pass "Restore with specific chain successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

test_restore_multiple_tables() {
    ((TESTS_RUN++))
    log_test "Restore multiple tables with relationships"

    local source_db="${TEST_DB_PREFIX}_source_multi"
    local restore_db="${TEST_DB_PREFIX}_restore_multi"
    local slot="test_slot_multi"
    local backup_dir="$TEST_DIR/multi"

    # Setup source database with multiple related tables
    create_test_db "$source_db"
    query_db "$source_db" "
        CREATE TABLE customers (
            id SERIAL PRIMARY KEY,
            name TEXT NOT NULL
        );
        CREATE TABLE orders (
            id SERIAL PRIMARY KEY,
            customer_id INTEGER REFERENCES customers(id),
            product TEXT NOT NULL
        );
    "
    query_db "$source_db" "INSERT INTO customers (name) VALUES ('Alice'), ('Bob');"
    query_db "$source_db" "INSERT INTO orders (customer_id, product) VALUES (1, 'Widget'), (2, 'Gadget');"

    # Initialize backup
    mkdir -p "$backup_dir"
    initialize_backup "$source_db" "$backup_dir" "$slot"

    # Restore
    if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
        DATABASES_TO_CLEANUP+=("$restore_db")

        # Verify customers
        local customer_count
        customer_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM customers;")
        if [[ "$customer_count" -ne 2 ]]; then
            log_fail "Expected 2 customers, got $customer_count"
            return 1
        fi

        # Verify orders
        local order_count
        order_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
        if [[ "$order_count" -ne 2 ]]; then
            log_fail "Expected 2 orders, got $order_count"
            return 1
        fi

        # Verify foreign key relationship
        local alice_orders
        alice_orders=$(query_db "$restore_db" "
            SELECT o.product FROM orders o
            JOIN customers c ON c.id = o.customer_id
            WHERE c.name = 'Alice';
        ")
        if [[ "$alice_orders" != "Widget" ]]; then
            log_fail "Foreign key relationship not preserved"
            return 1
        fi

        log_pass "Multiple table restore successful"
        return 0
    else
        log_fail "Restore command failed"
        return 1
    fi
}

#
# Cleanup
#

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

    # Stop any background pg_scribe processes
    pkill -f "pg_scribe.*--start" 2>/dev/null || true
    pkill -f "pg_recvlogical" 2>/dev/null || true
    sleep 1

    # Drop replication slots
    for slot_info in "${SLOTS_TO_CLEANUP[@]}"; do
        IFS=':' read -r dbname slot <<< "$slot_info"
        drop_replication_slot "$dbname" "$slot" 2>/dev/null || true
    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 --restore 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 ! psql -U "$PGUSER" -d postgres -c "CREATE EXTENSION IF NOT EXISTS wal2sql;" &>/dev/null; then
        echo "ERROR: wal2sql extension not available"
        echo "Build and install: 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_basic_restore_from_base_only || true
    test_restore_with_differentials || true
    test_restore_with_ddl_changes || true
    test_restore_specific_chain || true
    test_restore_to_existing_database || true
    test_restore_fails_if_db_exists_with_create || true
    test_restore_missing_backup_directory || true
    test_restore_multiple_tables || true
    test_restore_sequence_synchronization || true
    test_restore_no_sync_sequences || 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 "$@"