#!/usr/bin/env bash
set -euo pipefail

# Test harness for database backends
# Usage:
#   ./scripts/test-db-backends          # test SQLite only (default)
#   ./scripts/test-db-backends sqlite   # test SQLite
#   ./scripts/test-db-backends postgres # test PostgreSQL (starts Docker)
#   ./scripts/test-db-backends all      # test both

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*"; }
log_step() { echo -e "${BLUE}[STEP]${NC} $*"; }

BACKEND="${1:-sqlite}"
TEST_DB_PATH="/tmp/prefect-test-$$.db"
POSTGRES_CONTAINER="prefect-test-postgres"
POSTGRES_PORT=5433  # use non-standard port to avoid conflicts
POSTGRES_USER="prefect"
POSTGRES_PASSWORD="prefect"
POSTGRES_DB="prefect_test"
POSTGRES_URL="postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@localhost:${POSTGRES_PORT}/${POSTGRES_DB}"
SERVER_PID=""

cleanup() {
    # Kill server if running
    if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
        log_info "Stopping server (PID $SERVER_PID)..."
        kill "$SERVER_PID" 2>/dev/null || true
        wait "$SERVER_PID" 2>/dev/null || true
    fi

    # Clean up test database
    rm -f "$TEST_DB_PATH" 2>/dev/null || true
}
trap cleanup EXIT

start_postgres_docker() {
    log_step "Setting up PostgreSQL via Docker..."

    # Check if Docker is available
    if ! command -v docker &> /dev/null; then
        log_error "Docker not found. Install Docker to test PostgreSQL backend."
        return 1
    fi

    # Stop any existing container
    docker rm -f "$POSTGRES_CONTAINER" 2>/dev/null || true

    # Start PostgreSQL container
    log_info "Starting PostgreSQL container on port $POSTGRES_PORT..."
    if ! docker run -d \
        --name "$POSTGRES_CONTAINER" \
        -e POSTGRES_USER="$POSTGRES_USER" \
        -e POSTGRES_PASSWORD="$POSTGRES_PASSWORD" \
        -e POSTGRES_DB="$POSTGRES_DB" \
        -p "${POSTGRES_PORT}:5432" \
        postgres:16-alpine \
        > /dev/null 2>&1; then
        log_error "Failed to start PostgreSQL container. Is Docker running?"
        return 1
    fi

    # Wait for PostgreSQL to be ready
    log_info "Waiting for PostgreSQL to be ready..."
    local max_attempts=30
    local attempt=0
    while [[ $attempt -lt $max_attempts ]]; do
        if docker exec "$POSTGRES_CONTAINER" pg_isready -U "$POSTGRES_USER" &> /dev/null; then
            log_info "PostgreSQL is ready"
            return 0
        fi
        ((attempt++))
        sleep 1
    done

    log_error "PostgreSQL failed to start within ${max_attempts}s"
    return 1
}

stop_postgres_docker() {
    log_info "Stopping PostgreSQL container..."
    docker rm -f "$POSTGRES_CONTAINER" 2>/dev/null || true
}

# Sanity check: verify data in SQLite
sanity_check_sqlite() {
    local db_path="$1"
    local flow_id="$2"
    local flow_run_id="$3"

    log_step "Sanity checks: verifying SQLite data..."

    # Check flow exists
    local flow_count
    flow_count=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow WHERE id='$flow_id'")
    if [[ "$flow_count" != "1" ]]; then
        log_error "Flow not found in database (expected 1, got $flow_count)"
        return 1
    fi
    log_info "  flow exists: yes"

    # Check flow run exists
    local flow_run_count
    flow_run_count=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow_run WHERE id='$flow_run_id'")
    if [[ "$flow_run_count" != "1" ]]; then
        log_error "Flow run not found in database (expected 1, got $flow_run_count)"
        return 1
    fi
    log_info "  flow_run exists: yes"

    # Check state history
    local state_count
    state_count=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow_run_state WHERE flow_run_id='$flow_run_id'")
    if [[ "$state_count" -lt 1 ]]; then
        log_error "No state history for flow run (expected >= 1, got $state_count)"
        return 1
    fi
    log_info "  state history: $state_count states"

    # Check final state
    local final_state
    final_state=$(sqlite3 "$db_path" "SELECT type FROM flow_run_state WHERE flow_run_id='$flow_run_id' ORDER BY timestamp DESC LIMIT 1")
    log_info "  final state: $final_state"

    # Count total rows in key tables
    local total_flows total_runs total_states
    total_flows=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow")
    total_runs=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow_run")
    total_states=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM flow_run_state")
    log_info "  table counts: flows=$total_flows, runs=$total_runs, states=$total_states"

    return 0
}

# Sanity check: verify data in PostgreSQL
sanity_check_postgres() {
    local flow_id="$1"
    local flow_run_id="$2"

    log_step "Sanity checks: verifying PostgreSQL data..."

    # Check flow exists
    local flow_count
    flow_count=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow WHERE id='$flow_id'" 2>/dev/null | tr -d ' ')
    if [[ "$flow_count" != "1" ]]; then
        log_error "Flow not found in database (expected 1, got $flow_count)"
        return 1
    fi
    log_info "  flow exists: yes"

    # Check flow run exists
    local flow_run_count
    flow_run_count=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow_run WHERE id='$flow_run_id'" 2>/dev/null | tr -d ' ')
    if [[ "$flow_run_count" != "1" ]]; then
        log_error "Flow run not found in database (expected 1, got $flow_run_count)"
        return 1
    fi
    log_info "  flow_run exists: yes"

    # Check state history
    local state_count
    state_count=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow_run_state WHERE flow_run_id='$flow_run_id'" 2>/dev/null | tr -d ' ')
    if [[ "$state_count" -lt 1 ]]; then
        log_error "No state history for flow run (expected >= 1, got $state_count)"
        return 1
    fi
    log_info "  state history: $state_count states"

    # Check final state
    local final_state
    final_state=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT type FROM flow_run_state WHERE flow_run_id='$flow_run_id' ORDER BY timestamp DESC LIMIT 1" 2>/dev/null | tr -d ' ')
    log_info "  final state: $final_state"

    # Count total rows in key tables
    local total_flows total_runs total_states
    total_flows=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow" 2>/dev/null | tr -d ' ')
    total_runs=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow_run" 2>/dev/null | tr -d ' ')
    total_states=$(docker exec "$POSTGRES_CONTAINER" psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c \
        "SELECT COUNT(*) FROM flow_run_state" 2>/dev/null | tr -d ' ')
    log_info "  table counts: flows=$total_flows, runs=$total_runs, states=$total_states"

    return 0
}

test_sqlite() {
    log_info "=== Testing SQLite backend ==="

    export PREFECT_DATABASE_BACKEND=sqlite
    export PREFECT_DATABASE_PATH="$TEST_DB_PATH"

    # Remove any existing test database
    rm -f "$TEST_DB_PATH"

    # Build and run tests
    cd "$PROJECT_DIR"

    log_step "Building..."
    zig build 2>&1 || { log_error "Build failed"; return 1; }

    log_step "Running backend unit tests..."
    zig build test 2>&1 || { log_error "Unit tests failed"; return 1; }

    log_step "Starting server for integration tests..."
    ./zig-out/bin/prefect-server &
    SERVER_PID=$!
    sleep 2

    # Basic health check
    log_step "Health check..."
    if curl -s http://localhost:4200/api/health | grep -q "ok"; then
        log_info "Health check passed"
    else
        log_error "Health check failed"
        return 1
    fi

    # Test flow creation
    log_step "Testing flow creation..."
    FLOW_RESPONSE=$(curl -s -X POST http://localhost:4200/api/flows/ \
        -H "Content-Type: application/json" \
        -d '{"name": "test-flow-sqlite"}')

    if echo "$FLOW_RESPONSE" | grep -q '"id"'; then
        log_info "Flow creation passed"
        FLOW_ID=$(echo "$FLOW_RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    else
        log_error "Flow creation failed: $FLOW_RESPONSE"
        return 1
    fi

    # Test flow retrieval
    log_step "Testing flow retrieval..."
    FLOW_GET=$(curl -s "http://localhost:4200/api/flows/$FLOW_ID")
    if echo "$FLOW_GET" | grep -q "test-flow-sqlite"; then
        log_info "Flow retrieval passed"
    else
        log_error "Flow retrieval failed: $FLOW_GET"
        return 1
    fi

    # Test flow run creation
    log_step "Testing flow run creation..."
    FLOW_RUN_RESPONSE=$(curl -s -X POST http://localhost:4200/api/flow_runs/ \
        -H "Content-Type: application/json" \
        -d "{\"flow_id\": \"$FLOW_ID\", \"name\": \"test-run-1\"}")

    if echo "$FLOW_RUN_RESPONSE" | grep -q '"id"'; then
        log_info "Flow run creation passed"
        FLOW_RUN_ID=$(echo "$FLOW_RUN_RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    else
        log_error "Flow run creation failed: $FLOW_RUN_RESPONSE"
        return 1
    fi

    # Test state transition: PENDING -> RUNNING
    log_step "Testing state transition (PENDING -> RUNNING)..."
    STATE_RESPONSE=$(curl -s -X POST "http://localhost:4200/api/flow_runs/$FLOW_RUN_ID/set_state" \
        -H "Content-Type: application/json" \
        -d '{"state": {"type": "RUNNING", "name": "Running"}}')

    if echo "$STATE_RESPONSE" | grep -q "RUNNING"; then
        log_info "State transition to RUNNING passed"
    else
        log_error "State transition failed: $STATE_RESPONSE"
        return 1
    fi

    # Test state transition: RUNNING -> COMPLETED
    log_step "Testing state transition (RUNNING -> COMPLETED)..."
    STATE_RESPONSE=$(curl -s -X POST "http://localhost:4200/api/flow_runs/$FLOW_RUN_ID/set_state" \
        -H "Content-Type: application/json" \
        -d '{"state": {"type": "COMPLETED", "name": "Completed"}}')

    if echo "$STATE_RESPONSE" | grep -q "COMPLETED"; then
        log_info "State transition to COMPLETED passed"
    else
        log_error "State transition failed: $STATE_RESPONSE"
        return 1
    fi

    # Test block type creation
    log_step "Testing block type creation..."
    BLOCK_TYPE_RESPONSE=$(curl -s -X POST http://localhost:4200/api/block_types/ \
        -H "Content-Type: application/json" \
        -d '{"name": "TestBlock", "slug": "test-block"}')

    if echo "$BLOCK_TYPE_RESPONSE" | grep -q '"id"'; then
        log_info "Block type creation passed"
    else
        log_error "Block type creation failed: $BLOCK_TYPE_RESPONSE"
        return 1
    fi

    # Stop server before sanity checks
    kill "$SERVER_PID" 2>/dev/null || true
    wait "$SERVER_PID" 2>/dev/null || true
    SERVER_PID=""

    # Run sanity checks
    sanity_check_sqlite "$TEST_DB_PATH" "$FLOW_ID" "$FLOW_RUN_ID" || return 1

    log_info "=== SQLite backend tests PASSED ==="
    return 0
}

test_postgres() {
    log_info "=== Testing PostgreSQL backend ==="

    # Start PostgreSQL via Docker
    start_postgres_docker || return 1

    export PREFECT_DATABASE_BACKEND=postgres
    export PREFECT_DATABASE_URL="$POSTGRES_URL"

    cd "$PROJECT_DIR"

    log_step "Building..."
    zig build 2>&1 || { log_error "Build failed"; stop_postgres_docker; return 1; }

    log_step "Starting server..."
    ./zig-out/bin/prefect-server &
    SERVER_PID=$!
    sleep 3

    # Check if server started
    if ! kill -0 $SERVER_PID 2>/dev/null; then
        log_error "Server failed to start - check backend implementation"
        stop_postgres_docker
        return 1
    fi

    # Health check
    log_step "Health check..."
    if curl -s http://localhost:4200/api/health | grep -q "ok"; then
        log_info "Health check passed"
    else
        log_error "Health check failed"
        stop_postgres_docker
        return 1
    fi

    # Test flow creation
    log_step "Testing flow creation..."
    FLOW_RESPONSE=$(curl -s -X POST http://localhost:4200/api/flows/ \
        -H "Content-Type: application/json" \
        -d '{"name": "test-flow-postgres"}')

    if echo "$FLOW_RESPONSE" | grep -q '"id"'; then
        log_info "Flow creation passed"
        FLOW_ID=$(echo "$FLOW_RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    else
        log_error "Flow creation failed: $FLOW_RESPONSE"
        stop_postgres_docker
        return 1
    fi

    # Test flow retrieval
    log_step "Testing flow retrieval..."
    FLOW_GET=$(curl -s "http://localhost:4200/api/flows/$FLOW_ID")
    if echo "$FLOW_GET" | grep -q "test-flow-postgres"; then
        log_info "Flow retrieval passed"
    else
        log_error "Flow retrieval failed: $FLOW_GET"
        stop_postgres_docker
        return 1
    fi

    # Test flow run creation
    log_step "Testing flow run creation..."
    FLOW_RUN_RESPONSE=$(curl -s -X POST http://localhost:4200/api/flow_runs/ \
        -H "Content-Type: application/json" \
        -d "{\"flow_id\": \"$FLOW_ID\", \"name\": \"test-run-pg-1\"}")

    if echo "$FLOW_RUN_RESPONSE" | grep -q '"id"'; then
        log_info "Flow run creation passed"
        FLOW_RUN_ID=$(echo "$FLOW_RUN_RESPONSE" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)
    else
        log_error "Flow run creation failed: $FLOW_RUN_RESPONSE"
        stop_postgres_docker
        return 1
    fi

    # Test state transition: PENDING -> RUNNING
    log_step "Testing state transition (PENDING -> RUNNING)..."
    STATE_RESPONSE=$(curl -s -X POST "http://localhost:4200/api/flow_runs/$FLOW_RUN_ID/set_state" \
        -H "Content-Type: application/json" \
        -d '{"state": {"type": "RUNNING", "name": "Running"}}')

    if echo "$STATE_RESPONSE" | grep -q "RUNNING"; then
        log_info "State transition to RUNNING passed"
    else
        log_error "State transition failed: $STATE_RESPONSE"
        stop_postgres_docker
        return 1
    fi

    # Test state transition: RUNNING -> COMPLETED
    log_step "Testing state transition (RUNNING -> COMPLETED)..."
    STATE_RESPONSE=$(curl -s -X POST "http://localhost:4200/api/flow_runs/$FLOW_RUN_ID/set_state" \
        -H "Content-Type: application/json" \
        -d '{"state": {"type": "COMPLETED", "name": "Completed"}}')

    if echo "$STATE_RESPONSE" | grep -q "COMPLETED"; then
        log_info "State transition to COMPLETED passed"
    else
        log_error "State transition failed: $STATE_RESPONSE"
        stop_postgres_docker
        return 1
    fi

    # Test block type creation
    log_step "Testing block type creation..."
    BLOCK_TYPE_RESPONSE=$(curl -s -X POST http://localhost:4200/api/block_types/ \
        -H "Content-Type: application/json" \
        -d '{"name": "TestBlockPG", "slug": "test-block-pg"}')

    if echo "$BLOCK_TYPE_RESPONSE" | grep -q '"id"'; then
        log_info "Block type creation passed"
    else
        log_error "Block type creation failed: $BLOCK_TYPE_RESPONSE"
        stop_postgres_docker
        return 1
    fi

    # Stop server before sanity checks
    kill "$SERVER_PID" 2>/dev/null || true
    wait "$SERVER_PID" 2>/dev/null || true
    SERVER_PID=""

    # Run sanity checks
    sanity_check_postgres "$FLOW_ID" "$FLOW_RUN_ID" || {
        stop_postgres_docker
        return 1
    }

    # Clean up Docker
    stop_postgres_docker

    log_info "=== PostgreSQL backend tests PASSED ==="
    return 0
}

main() {
    case "$BACKEND" in
        sqlite)
            test_sqlite
            ;;
        postgres|postgresql)
            test_postgres
            ;;
        all)
            log_info "Running all backend tests..."
            echo ""
            test_sqlite || exit 1
            echo ""
            test_postgres || exit 1
            echo ""
            log_info "=== ALL BACKEND TESTS PASSED ==="
            ;;
        *)
            log_error "Unknown backend: $BACKEND"
            echo "Usage: $0 [sqlite|postgres|all]"
            exit 1
            ;;
    esac
}

main