#!/usr/bin/env bash
#
# Test suite for pg_scribe --stop command
#
# This test suite:
# - Creates temporary test databases
# - Tests various --stop 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
}

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 a backup directory (creates replication slot and initial backups)
init_backup_system() {
    local dbname="$1"
    local backup_dir="$2"
    local slot="$3"

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

# Start streaming in the background
start_streaming_background() {
    local dbname="$1"
    local backup_dir="$2"
    local slot="$3"

    "$PG_SCRIBE" --start -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
    local pid=$!

    # Wait for pidfile to appear (instead of fixed sleep)
    local pidfile="$backup_dir/.pg_scribe.pid"
    local timeout=10
    local waited=0

    while [[ ! -f "$pidfile" ]] && [[ $waited -lt $timeout ]]; do
        sleep 0.5
        waited=$((waited + 1))
    done

    # Verify process is still running
    if ! kill -0 "$pid" 2>/dev/null; then
        return 1
    fi

    # Verify pidfile was created
    if [[ ! -f "$pidfile" ]]; then
        return 1
    fi

    echo "$pid"
}

#
# Test cases
#

test_stop_requires_backup_dir() {
    ((TESTS_RUN++))
    log_test "Stop requires backup directory"

    local exit_code=0
    "$PG_SCRIBE" --stop &>/dev/null || exit_code=$?

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

test_stop_nonexistent_directory() {
    ((TESTS_RUN++))
    log_test "Stop fails with nonexistent directory"

    local backup_dir="$TEST_DIR/nonexistent"

    local exit_code=0
    "$PG_SCRIBE" --stop -f "$backup_dir" -U "$PGUSER" &>/dev/null || exit_code=$?

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

test_stop_no_active_process() {
    ((TESTS_RUN++))
    log_test "Stop fails when no active process"

    local dbname="${TEST_DB_PREFIX}_noprocess"
    local backup_dir="$TEST_DIR/noprocess"
    local slot="test_slot_noprocess"

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

    # Try to stop (should fail - no process running)
    local exit_code=0
    "$PG_SCRIBE" --stop -f "$backup_dir" -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 1 ]]; then
        log_pass "Correctly fails when no process running"
        return 0
    else
        log_fail "Expected exit code 1, got $exit_code"
        return 1
    fi
}

test_stop_active_process() {
    ((TESTS_RUN++))
    log_test "Stop successfully stops active streaming"

    local dbname="${TEST_DB_PREFIX}_active"
    local backup_dir="$TEST_DIR/active"
    local slot="test_slot_active"

    # Setup - initialize and start streaming
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    init_backup_system "$dbname" "$backup_dir" "$slot"

    # Start streaming in background
    local pid
    pid=$(start_streaming_background "$dbname" "$backup_dir" "$slot")

    if [[ -z "$pid" ]]; then
        log_fail "Failed to start streaming"
        return 1
    fi

    # Stop the process
    if "$PG_SCRIBE" --stop -f "$backup_dir" -U "$PGUSER" &>/dev/null; then
        # Verify process is actually stopped
        sleep 1
        if kill -0 "$pid" 2>/dev/null; then
            log_fail "Process still running after stop"
            kill -KILL "$pid" 2>/dev/null || true
            return 1
        fi

        # Verify pidfile was removed
        local pidfile="$backup_dir/.pg_scribe.pid"
        if [[ -f "$pidfile" ]]; then
            log_fail "Pidfile not removed"
            return 1
        fi

        log_pass "Successfully stopped active streaming"
        return 0
    else
        log_fail "Stop command failed"
        kill -TERM "$pid" 2>/dev/null || true
        return 1
    fi
}

test_stop_stale_pidfile() {
    ((TESTS_RUN++))
    log_test "Stop handles stale pidfile gracefully"

    local dbname="${TEST_DB_PREFIX}_stale"
    local backup_dir="$TEST_DIR/stale"
    local slot="test_slot_stale"

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

    # Create stale pidfile (with non-existent PID)
    local pidfile="$backup_dir/.pg_scribe.pid"
    echo "99999" > "$pidfile"

    # Stop should handle this gracefully (exit 0 and remove stale pidfile)
    if "$PG_SCRIBE" --stop -f "$backup_dir" -U "$PGUSER" &>/dev/null; then
        # Verify pidfile was removed
        if [[ -f "$pidfile" ]]; then
            log_fail "Stale pidfile not removed"
            return 1
        fi

        log_pass "Stale pidfile handled gracefully"
        return 0
    else
        log_fail "Stop failed on stale pidfile"
        return 1
    fi
}

#
# Cleanup
#

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

    # Kill any remaining pg_recvlogical processes
    pkill -f "pg_recvlogical.*test_slot" 2>/dev/null || true
    sleep 1

    # Drop replication slots
    for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
        for slot in test_slot_noprocess test_slot_active test_slot_stale; 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 --stop 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_stop_requires_backup_dir || true
    test_stop_nonexistent_directory || true
    test_stop_no_active_process || true
    test_stop_active_process || true
    test_stop_stale_pidfile || 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 "$@"