#!/usr/bin/env bash
#
# Test suite for pg_scribe --init command
#
# This test suite:
# - Creates temporary test databases
# - Tests various --init 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_$$"
TEST_DB_PREFIX="pg_scribe_test_$$"
PGUSER="${PGUSER:-postgres}"

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

# Cleanup tracking
DATABASES_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
}

check_slot_exists() {
    local dbname="$1"
    local slot="$2"
    local exists
    exists=$(query_db "$dbname" "
        SELECT 1 FROM pg_replication_slots WHERE slot_name = '$slot';
    ")
    [[ -n "$exists" ]]
}

check_extension_exists() {
    local dbname="$1"
    local extension="$2"
    local exists
    exists=$(query_db "$dbname" "
        SELECT 1 FROM pg_extension WHERE extname = '$extension';
    ")
    [[ -n "$exists" ]]
}

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()
        );
    "
}

create_table_without_pk() {
    local dbname="$1"
    local table="$2"
    query_db "$dbname" "
        CREATE TABLE $table (
            id INTEGER,
            name TEXT
        );
    "
}

#
# Test cases
#

test_basic_init_success() {
    ((TESTS_RUN++))
    log_test "Basic --init success"

    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"
    mkdir -p "$backup_dir"

    # Run init
    if "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null; then
        # Verify slot created
        if ! check_slot_exists "$dbname" "$slot"; then
            log_fail "Replication slot not created"
            return 1
        fi

        # Verify extension installed
        if ! check_extension_exists "$dbname" "wal2sql"; then
            log_fail "wal2sql extension not installed"
            return 1
        fi

        # Verify chain directory created
        local chain_dirs=("$backup_dir"/chain-*)
        if [[ ! -d "${chain_dirs[0]}" ]]; then
            log_fail "Chain directory not created"
            return 1
        fi

        local chain_dir="${chain_dirs[0]}"

        # Verify files in chain directory
        if [[ ! -f "$chain_dir/base.sql" ]]; then
            log_fail "Base backup file not created in chain"
            return 1
        fi

        if [[ ! -f "$chain_dir/globals.sql" ]]; then
            log_fail "Globals backup file not created in chain"
            return 1
        fi

        if [[ ! -f "$chain_dir/metadata.json" ]]; then
            log_fail "Metadata file not created in chain"
            return 1
        fi

        # Verify pidfile is NOT created (only created during --start)
        if [[ -f "$backup_dir/.pg_scribe.pid" ]]; then
            log_fail "Pidfile should not be created by --init"
            return 1
        fi

        # Verify metadata content (JSON format)
        if ! grep -q "\"database\": \"$dbname\"" "$chain_dir/metadata.json"; then
            log_fail "Metadata missing database name"
            return 1
        fi

        if ! grep -q "\"replication_slot\": \"$slot\"" "$chain_dir/metadata.json"; then
            log_fail "Metadata missing slot name"
            return 1
        fi

        log_pass "Basic init successful"
        return 0
    else
        log_fail "Init command failed"
        return 1
    fi
}

test_init_validation_failure() {
    ((TESTS_RUN++))
    log_test "Init validation failure (table without replica identity)"

    local dbname="${TEST_DB_PREFIX}_nopk"
    local slot="test_slot_nopk"
    local backup_dir="$TEST_DIR/nopk"

    # Setup - create table WITHOUT primary key
    create_test_db "$dbname"
    create_table_without_pk "$dbname" "bad_table"
    mkdir -p "$backup_dir"

    # Run init - should fail with exit code 5
    local exit_code=0
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 5 ]]; then
        # Verify slot was NOT created
        if check_slot_exists "$dbname" "$slot"; then
            log_fail "Replication slot should not be created on validation failure"
            return 1
        fi

        log_pass "Validation failure detected correctly"
        return 0
    else
        log_fail "Expected exit code 5, got $exit_code"
        return 1
    fi
}

test_init_force_flag() {
    ((TESTS_RUN++))
    log_test "Init with --force flag (bypass validation)"

    local dbname="${TEST_DB_PREFIX}_force"
    local slot="test_slot_force"
    local backup_dir="$TEST_DIR/force"

    # Setup - create table WITHOUT primary key
    create_test_db "$dbname"
    create_table_without_pk "$dbname" "bad_table"
    mkdir -p "$backup_dir"

    # Run init with --force - should succeed despite validation failure
    if "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" --force &>/dev/null; then
        # Verify slot was created
        if ! check_slot_exists "$dbname" "$slot"; then
            log_fail "Replication slot should be created with --force"
            return 1
        fi

        log_pass "Force flag bypassed validation"
        return 0
    else
        log_fail "Init with --force should succeed"
        return 1
    fi
}

test_init_non_idempotency() {
    ((TESTS_RUN++))
    log_test "Init refuses to run on already-initialized directory"

    local dbname="${TEST_DB_PREFIX}_nonidempotent"
    local slot="test_slot_nonidempotent"
    local backup_dir="$TEST_DIR/nonidempotent"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    mkdir -p "$backup_dir"

    # Run init first time
    if ! "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null; then
        log_fail "First init failed"
        return 1
    fi

    # Run init second time - should FAIL with validation error
    local exit_code=0
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 5 ]]; then
        # Verify slot still exists from first init
        if ! check_slot_exists "$dbname" "$slot"; then
            log_fail "Replication slot should still exist from first init"
            return 1
        fi

        # Verify only 1 chain directory (from first init)
        local chain_count
        chain_count=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | wc -l)
        if [[ $chain_count -ne 1 ]]; then
            log_fail "Expected 1 chain directory, got $chain_count"
            return 1
        fi

        log_pass "Init correctly refuses to reinitialize"
        return 0
    else
        log_fail "Expected exit code 5 (validation error), got $exit_code"
        return 1
    fi
}

test_if_not_exists_flag() {
    ((TESTS_RUN++))
    log_test "Init with --if-not-exists flag (idempotent)"

    local dbname="${TEST_DB_PREFIX}_ifnotexists"
    local slot="test_slot_ifnotexists"
    local backup_dir="$TEST_DIR/ifnotexists"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    mkdir -p "$backup_dir"

    # Run init first time
    if ! "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null; then
        log_fail "First init failed"
        return 1
    fi

    # Verify slot was created
    if ! check_slot_exists "$dbname" "$slot"; then
        log_fail "Replication slot not created on first init"
        return 1
    fi

    # Run init second time with --if-not-exists - should SUCCEED with exit 0
    local exit_code=0
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" --if-not-exists &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 0 ]]; then
        # Verify slot still exists from first init
        if ! check_slot_exists "$dbname" "$slot"; then
            log_fail "Replication slot should still exist"
            return 1
        fi

        # Verify only 1 chain directory (from first init, second init should skip)
        local chain_count
        chain_count=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | wc -l)
        if [[ $chain_count -ne 1 ]]; then
            log_fail "Expected 1 chain directory, got $chain_count"
            return 1
        fi

        log_pass "--if-not-exists flag makes init idempotent"
        return 0
    else
        log_fail "Expected exit code 0, got $exit_code"
        return 1
    fi
}

test_init_multiple_tables() {
    ((TESTS_RUN++))
    log_test "Init with multiple tables"

    local dbname="${TEST_DB_PREFIX}_multi"
    local slot="test_slot_multi"
    local backup_dir="$TEST_DIR/multi"

    # Setup - create multiple tables
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    create_table_with_pk "$dbname" "orders"
    create_table_with_pk "$dbname" "products"

    # Insert some data
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Alice'), ('Bob');"
    query_db "$dbname" "INSERT INTO orders (name) VALUES ('Order1');"
    query_db "$dbname" "INSERT INTO products (name) VALUES ('Widget');"

    mkdir -p "$backup_dir"

    # Run init
    if "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null; then
        # Verify base backup contains all tables
        local chain_dirs=("$backup_dir"/chain-*)
        local base_file="${chain_dirs[0]}/base.sql"

        if ! grep -q "CREATE TABLE public.users" "$base_file"; then
            log_fail "Base backup missing users table"
            return 1
        fi

        if ! grep -q "CREATE TABLE public.orders" "$base_file"; then
            log_fail "Base backup missing orders table"
            return 1
        fi

        if ! grep -q "CREATE TABLE public.products" "$base_file"; then
            log_fail "Base backup missing products table"
            return 1
        fi

        log_pass "Multiple tables backed up successfully"
        return 0
    else
        log_fail "Init failed"
        return 1
    fi
}

test_init_with_unlogged_table() {
    ((TESTS_RUN++))
    log_test "Init with unlogged table (warning but success)"

    local dbname="${TEST_DB_PREFIX}_unlogged"
    local slot="test_slot_unlogged"
    local backup_dir="$TEST_DIR/unlogged"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "normal_table"
    query_db "$dbname" "CREATE UNLOGGED TABLE unlogged_table (id SERIAL PRIMARY KEY, data TEXT);"
    mkdir -p "$backup_dir"

    # Run init - should succeed with warning (exit code 0 or 10)
    local output_file="$TEST_DIR/unlogged_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Exit code 0 (success) or 10 (warning) are both acceptable
    if [[ $exit_code -eq 0 || $exit_code -eq 10 ]]; then
        # Check for warning message
        if ! grep -i "unlogged" "$output_file"; then
            log_fail "Expected warning about unlogged table"
            return 1
        fi

        # Verify backup was created
        local chain_dirs=("$backup_dir"/chain-*)
        if [[ ! -f "${chain_dirs[0]}/base.sql" ]]; then
            log_fail "Base backup file not created"
            return 1
        fi

        log_pass "Unlogged table warning shown correctly"
        return 0
    else
        log_fail "Init failed with exit code $exit_code"
        return 1
    fi
}

test_backup_content_validity() {
    ((TESTS_RUN++))
    log_test "Verify backup SQL is valid"

    local dbname="${TEST_DB_PREFIX}_valid"
    local slot="test_slot_valid"
    local backup_dir="$TEST_DIR/valid"
    local restore_dbname="${TEST_DB_PREFIX}_restored"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "test_data"
    query_db "$dbname" "INSERT INTO test_data (name) VALUES ('Test Row 1'), ('Test Row 2');"
    mkdir -p "$backup_dir"

    # Run init
    if ! "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null; then
        log_fail "Init failed"
        return 1
    fi

    # Try to restore the backup to a new database
    create_test_db "$restore_dbname"

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

    local base_file="$chain_dir/base.sql"
    local globals_file="$chain_dir/globals.sql"

    # Apply globals (roles, etc.)
    if ! psql -U "$PGUSER" -d postgres -f "$globals_file" &>/dev/null; then
        log_fail "Failed to restore globals"
        return 1
    fi

    # Apply base backup
    if ! psql -U "$PGUSER" -d "$restore_dbname" -f "$base_file" &>/dev/null; then
        log_fail "Failed to restore base backup"
        return 1
    fi

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

    log_pass "Backup SQL is valid and restorable"
    return 0
}

#
# Cleanup
#

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

    # Drop replication slots
    for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
        # Try to drop any slots for this database
        for slot in test_slot_basic test_slot_nopk test_slot_force test_slot_nonidempotent test_slot_ifnotexists test_slot_multi test_slot_unlogged test_slot_valid; 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 --init 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

    # 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_init_success || true
    test_init_validation_failure || true
    test_init_force_flag || true
    test_init_non_idempotency || true
    test_if_not_exists_flag || true
    test_init_multiple_tables || true
    test_init_with_unlogged_table || true
    test_backup_content_validity || 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 "$@"