#!/usr/bin/env bash
#
# Test suite for pg_scribe --status command
#
# This test suite:
# - Creates temporary test databases
# - Tests various --status scenarios
# - Verifies status reporting and health checks
# - 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_status_$$"
TEST_DB_PREFIX="pg_scribe_test_status_$$"
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
}

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" ]]
}

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 (exit on failure - this is critical for tests)
    local init_output
    init_output=$("$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" 2>&1)
    local init_exit=$?

    if [[ $init_exit -ne 0 && $init_exit -ne 10 ]]; then
        log_info "ERROR: Failed to initialize backup system for $dbname (exit=$init_exit)"
        echo "$init_output" >&2
        return 1
    fi
    return 0
}

#
# Test cases
#

test_status_without_slot() {
    ((TESTS_RUN++))
    log_test "Status check when slot doesn't exist (should fail)"

    local dbname="${TEST_DB_PREFIX}_noslot"
    local slot="test_slot_noslot"

    # Setup - create db but DON'T initialize
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"

    # Try to check status - should fail with exit code 3 (slot error)
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -U "$PGUSER" &>/dev/null || exit_code=$?

    if [[ $exit_code -eq 3 ]]; then
        log_pass "Correctly failed with slot error"
        return 0
    else
        log_fail "Expected exit code 3, got $exit_code"
        return 1
    fi
}

test_status_basic_healthy() {
    ((TESTS_RUN++))
    log_test "Basic status check (without backup directory)"

    local dbname="${TEST_DB_PREFIX}_healthy"
    local slot="test_slot_healthy"
    local backup_dir="$TEST_DIR/healthy"

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

    # Check status WITHOUT backup directory - should succeed
    # (warnings about inactive slot, but exit code 10 is acceptable)
    local output_file="$TEST_DIR/healthy_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Exit code 0 or 10 (warning about inactive slot) are acceptable
    if [[ $exit_code -ne 0 && $exit_code -ne 10 ]]; then
        log_fail "Status check failed with unexpected exit code $exit_code"
        return 1
    fi

    # Verify output contains expected information
    if ! grep -q "Slot Name:" "$output_file"; then
        log_fail "Output missing slot name"
        return 1
    fi

    if ! grep -q "Active:" "$output_file"; then
        log_fail "Output missing active status"
        return 1
    fi

    if ! grep -q "Restart Lag:" "$output_file"; then
        log_fail "Output missing lag information"
        return 1
    fi

    log_pass "Status check successful"
    return 0
}

test_status_with_backup_directory() {
    ((TESTS_RUN++))
    log_test "Status check with backup directory analysis"

    local dbname="${TEST_DB_PREFIX}_with_dir"
    local slot="test_slot_with_dir"
    local backup_dir="$TEST_DIR/with_dir"

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

    # Check status with backup directory
    local output_file="$TEST_DIR/with_dir_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -f "$backup_dir" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Exit code 0 or 10 (warning) are both acceptable here
    # (warning because no incremental backups exist yet)
    if [[ $exit_code -ne 0 && $exit_code -ne 10 ]]; then
        log_fail "Status check failed with unexpected exit code $exit_code"
        return 1
    fi

    # Verify backup directory analysis appears in output
    if ! grep -q "Backup Directory Analysis" "$output_file"; then
        log_fail "Output missing backup directory analysis"
        return 1
    fi

    if ! grep -q "Base Backups:" "$output_file"; then
        log_fail "Output missing base backup count"
        return 1
    fi

    if ! grep -q "Metadata File:" "$output_file"; then
        log_fail "Output missing metadata file status"
        return 1
    fi

    log_pass "Backup directory analysis works correctly"
    return 0
}

test_status_inactive_slot() {
    ((TESTS_RUN++))
    log_test "Status check with inactive slot (should warn)"

    local dbname="${TEST_DB_PREFIX}_inactive"
    local slot="test_slot_inactive"
    local backup_dir="$TEST_DIR/inactive"

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

    # Slot exists but is inactive (no pg_recvlogical running)
    # Check status - should succeed but with warning
    local output_file="$TEST_DIR/inactive_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Should exit with warning code (10)
    if [[ $exit_code -ne 10 ]]; then
        log_fail "Expected exit code 10 (warning), got $exit_code"
        return 1
    fi

    # Should mention inactive status
    if ! grep -qi "not active" "$output_file"; then
        log_fail "Output should mention inactive slot"
        return 1
    fi

    log_pass "Inactive slot warning works correctly"
    return 0
}

test_status_with_active_streaming() {
    ((TESTS_RUN++))
    log_test "Status check with active streaming"

    local dbname="${TEST_DB_PREFIX}_streaming"
    local slot="test_slot_streaming"
    local backup_dir="$TEST_DIR/streaming"
    local output_file="$TEST_DIR/streaming_incremental.sql"

    # 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 "$output_file" -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 activate
    sleep 2

    # Check status - slot should now be active
    local status_output="$TEST_DIR/streaming_status.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -U "$PGUSER" &>"$status_output" || exit_code=$?

    # 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 status showed active slot
    if ! grep -q "Active:.*Yes" "$status_output"; then
        log_fail "Slot should be shown as active"
        return 1
    fi

    log_pass "Active slot detected correctly"
    return 0
}

test_status_with_incremental_backups() {
    ((TESTS_RUN++))
    log_test "Status check with incremental backup files"

    local dbname="${TEST_DB_PREFIX}_with_inc"
    local slot="test_slot_with_inc"
    local backup_dir="$TEST_DIR/with_inc"
    local output_file="$backup_dir/incremental.sql"

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

    # Start streaming to create incremental backup
    "$PG_SCRIBE" --start -d "$dbname" -f "$output_file" -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 to generate incremental data
    query_db "$dbname" "INSERT INTO users (name) VALUES ('Test User');"

    # 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

    # Check status with backup directory
    local status_output="$TEST_DIR/with_inc_status.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -f "$backup_dir" -U "$PGUSER" &>"$status_output" || exit_code=$?

    # Should succeed (exit code 0 or 10 for warnings like inactive slot)
    if [[ $exit_code -ne 0 && $exit_code -ne 10 ]]; then
        log_fail "Status check failed with unexpected exit code $exit_code"
        return 1
    fi

    # Verify incremental files are reported
    if ! grep -q "Incremental Files:" "$status_output"; then
        log_fail "Output missing incremental file count"
        return 1
    fi

    # Should show at least 1 incremental file
    if grep -q "Incremental Files:.*0" "$status_output"; then
        log_fail "Should report at least 1 incremental file"
        return 1
    fi

    log_pass "Incremental backup files reported correctly"
    return 0
}

test_status_missing_metadata() {
    ((TESTS_RUN++))
    log_test "Status check with missing metadata file (should warn)"

    local dbname="${TEST_DB_PREFIX}_no_meta"
    local slot="test_slot_no_meta"
    local backup_dir="$TEST_DIR/no_meta"

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

    # Remove metadata file
    rm -f "$backup_dir/pg_scribe_metadata.txt"

    # Check status
    local output_file="$TEST_DIR/no_meta_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -f "$backup_dir" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Should exit with warning (10)
    if [[ $exit_code -ne 10 ]]; then
        log_fail "Expected exit code 10 (warning), got $exit_code"
        return 1
    fi

    # Should warn about missing metadata
    if ! grep -qi "metadata.*not found" "$output_file"; then
        log_fail "Should warn about missing metadata file"
        return 1
    fi

    log_pass "Missing metadata warning works correctly"
    return 0
}

test_status_nonexistent_backup_dir() {
    ((TESTS_RUN++))
    log_test "Status check with nonexistent backup directory"

    local dbname="${TEST_DB_PREFIX}_no_dir"
    local slot="test_slot_no_dir"
    local backup_dir="$TEST_DIR/no_dir"
    local fake_dir="$TEST_DIR/does_not_exist"

    # Setup - initialize but check status with different directory
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"
    initialize_backup_system "$dbname" "$slot" "$backup_dir"

    # Check status with nonexistent directory
    local output_file="$TEST_DIR/no_dir_output.txt"
    local exit_code=0
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -f "$fake_dir" -U "$PGUSER" &>"$output_file" || exit_code=$?

    # Should exit with warning (10)
    if [[ $exit_code -ne 10 ]]; then
        log_fail "Expected exit code 10 (warning), got $exit_code"
        return 1
    fi

    # Should warn about nonexistent directory
    if ! grep -qi "does not exist" "$output_file"; then
        log_fail "Should warn about nonexistent directory"
        return 1
    fi

    log_pass "Nonexistent directory warning works correctly"
    return 0
}

test_status_lag_display() {
    ((TESTS_RUN++))
    log_test "Status displays lag information correctly"

    local dbname="${TEST_DB_PREFIX}_lag"
    local slot="test_slot_lag"
    local backup_dir="$TEST_DIR/lag"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"

    if ! initialize_backup_system "$dbname" "$slot" "$backup_dir"; then
        log_fail "Failed to initialize backup system"
        return 1
    fi

    # Check status
    local output_file="$TEST_DIR/lag_output.txt"
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -U "$PGUSER" &>"$output_file" || true

    # Verify lag is displayed - use case-insensitive grep for ANSI color codes
    if ! grep -i "restart.*lag" "$output_file"; then
        log_fail "Output missing restart lag"
        cat "$output_file" >&2
        return 1
    fi

    if ! grep -i "confirmed.*lag" "$output_file"; then
        log_fail "Output missing confirmed lag"
        cat "$output_file" >&2
        return 1
    fi

    # Verify lag values are in MB
    if ! grep "MB" "$output_file"; then
        log_fail "Lag should be displayed in MB"
        cat "$output_file" >&2
        return 1
    fi

    log_pass "Lag information displayed correctly"
    return 0
}

test_status_multiple_base_backups() {
    ((TESTS_RUN++))
    log_test "Status with multiple base backups shows latest"

    local dbname="${TEST_DB_PREFIX}_multi_base"
    local slot="test_slot_multi_base"
    local backup_dir="$TEST_DIR/multi_base"

    # Setup
    create_test_db "$dbname"
    create_table_with_pk "$dbname" "users"

    if ! initialize_backup_system "$dbname" "$slot" "$backup_dir"; then
        log_fail "Failed to initialize backup system"
        return 1
    fi

    # Create additional base backup files manually
    sleep 1
    touch "$backup_dir/base-20250101-000000.sql"
    sleep 1
    touch "$backup_dir/base-20250102-000000.sql"

    # Check status
    local output_file="$TEST_DIR/multi_base_output.txt"
    "$PG_SCRIBE" --status -d "$dbname" -S "$slot" -f "$backup_dir" -U "$PGUSER" &>"$output_file" || true

    # Count base backups - should be 3
    local count
    count=$(find "$backup_dir" -maxdepth 1 -name 'base-*.sql' 2>/dev/null | wc -l)

    # Should report multiple base backups (at least 3)
    if ! grep -q "Base Backups:" "$output_file"; then
        log_fail "Output missing base backup count"
        cat "$output_file" >&2
        return 1
    fi

    # Check that we have at least 3 base backups
    if [[ $count -lt 3 ]]; then
        log_fail "Expected at least 3 base backups, found $count"
        ls -la "$backup_dir" >&2
        return 1
    fi

    # Should show latest base backup
    if ! grep -q "Latest Base:" "$output_file"; then
        log_fail "Should show latest base backup"
        cat "$output_file" >&2
        return 1
    fi

    log_pass "Multiple base backups handled 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_noslot test_slot_healthy test_slot_with_dir test_slot_inactive test_slot_streaming test_slot_with_inc test_slot_no_meta test_slot_no_dir test_slot_lag test_slot_multi_base; 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 --status 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_status_without_slot || true
    test_status_basic_healthy || true
    test_status_with_backup_directory || true
    test_status_inactive_slot || true
    test_status_with_active_streaming || true
    test_status_with_incremental_backups || true
    test_status_missing_metadata || true
    test_status_nonexistent_backup_dir || true
    test_status_lag_display || true
    test_status_multiple_base_backups || 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 "$@"