#!/usr/bin/env bash
# ci/scripts/run-openbao.sh — OpenBao multi-pod integration tests
#
# Starts a 3-node OpenBao Raft cluster (bao-1/2/3) plus a bootstrap
# container (bao-init) and a PostgreSQL instance that authenticates via
# AppRole (pg-bao).  Validates all v1.3+ KMS features against a real
# (non-mock) Vault-API-compatible KMS.
#
# Test plan:
#   1  OpenBao cluster is sealed=false and has an active leader
#   2  AppRole credentials were generated by bao-init
#   3  PostgreSQL + pg_vault_tde started successfully on pg-bao
#   4  pg_vault_tde_vault_fetch_dek() acquires a real DEK from OpenBao
#   5  Encryption/decryption round-trip (encrypted_heap INSERT → SELECT)
#   6  pg_vault_tde_health_check() shows vault_reachable=true + approle auth
#   7  KEK wrapping: wrapped_dek file written to PGDATA
#   8  KEK unwrap on simulated restart (pg_vault_tde_vault_fetch_dek re-runs)
#   9  pg_vault_tde_vault_rewrap_dek() advances key version in OpenBao
#  10  BGW token renewal: pg_vault_tde.bgw_enabled=on — token auto-renewed
#  11  Node bao-2 connectivity: bao-2 API returns 200 health (follower up)
#  12  Node bao-3 connectivity: bao-3 API returns 200 health (follower up)
#
# Exit code: 0 on success, 5 on failure.
#
# Environment variables:
#   BAO_ROOT_TOKEN     — Root token (default: bao-root-token)
#   BAO_MOUNT          — Transit mount path (default: transit)
#   BAO_KEY_NAME       — Transit key name (default: pg-tde-dek)
#   BAO_PORT_1         — Host port for bao-1 (default: 18200)
#   BAO_PORT_2         — Host port for bao-2 (default: 18201)
#   BAO_PORT_3         — Host port for bao-3 (default: 18202)
#   PG_BAO_PORT        — Host port for pg-bao (default: 18432)
#   OPENBAO_STARTUP_TIMEOUT — seconds to wait for cluster (default: 60)
#
# Copyright (c) 2026 Miriade S.r.l. — PostgreSQL License (BSD)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=ci/scripts/lib.sh
source "$SCRIPT_DIR/lib.sh"

COMPOSE_FILE="$CI_DIR/compose-openbao.yml"
BAO_ROOT_TOKEN="${BAO_ROOT_TOKEN:-bao-root-token}"
BAO_MOUNT="${BAO_MOUNT:-transit}"
BAO_KEY_NAME="${BAO_KEY_NAME:-pg-tde-dek}"
BAO_PORT_1="${BAO_PORT_1:-18200}"
BAO_PORT_2="${BAO_PORT_2:-18201}"
BAO_PORT_3="${BAO_PORT_3:-18202}"
PG_BAO_PORT="${PG_BAO_PORT:-18432}"
STARTUP_TIMEOUT="${OPENBAO_STARTUP_TIMEOUT:-60}"

TESTS_PASSED=0
TESTS_FAILED=0

cleanup() {
    log_info "Tearing down OpenBao compose services ..."
    cd "$CI_DIR"
    $COMPOSE_CMD -f "$COMPOSE_FILE" down -v --remove-orphans 2>/dev/null || true
}
trap cleanup EXIT

# ---------------------------------------------------------------------------
# Test assertion helpers
# ---------------------------------------------------------------------------

pass() {
    TESTS_PASSED=$(( TESTS_PASSED + 1 ))
    log_ok "  PASS: $*"
}

fail() {
    TESTS_FAILED=$(( TESTS_FAILED + 1 ))
    log_error "  FAIL: $*"
}

# run_bao_test TEST_NUM DESCRIPTION SQL — executes SQL on pg-bao, asserts success
run_bao_test() {
    local num="$1" desc="$2" sql="$3"
    log_info "Test $num: $desc ..."
    if $RT exec -u postgres pg-tde-bao \
            psql -U postgres -At -c "$sql" > /dev/null 2>&1; then
        pass "Test $num: $desc"
    else
        fail "Test $num: $desc"
    fi
}

# bao_api_ok TEST_NUM DESCRIPTION PORT PATH — checks GET returns 200
bao_api_ok() {
    local num="$1" desc="$2" port="$3" path="$4"
    log_info "Test $num: $desc ..."
    if $RT exec bao-1 wget -qO- "http://localhost:${port}${path}" \
            2>/dev/null | grep -q '"sealed":false'; then
        pass "Test $num: $desc"
    else
        # Fall back to host-side check via published port
        if wget -qO- "http://127.0.0.1:${port}${path}" \
                2>/dev/null | grep -q '"sealed":false'; then
            pass "Test $num: $desc"
        else
            fail "Test $num: $desc"
        fi
    fi
}

# ---------------------------------------------------------------------------
# Build & start
# ---------------------------------------------------------------------------
log_stage "OPENBAO MULTI-POD INTEGRATION TESTS"

log_info "Building pg-tde-test image ..."
build_pg_test_image

log_info "Starting OpenBao cluster (bao-1, bao-2, bao-3 + bao-init + pg-bao) ..."
cd "$CI_DIR"
$COMPOSE_CMD -f "$COMPOSE_FILE" up -d bao-1 bao-2 bao-3

# Wait for the Raft leader (bao-1) to be healthy and unsealed
log_info "Waiting for OpenBao Raft leader (bao-1) ..."
BAO_1_URL="http://127.0.0.1:${BAO_PORT_1}"
for i in $(seq 1 "$STARTUP_TIMEOUT"); do
    status=$(wget -qO- "${BAO_1_URL}/v1/sys/health" 2>/dev/null || true)
    if echo "$status" | grep -q '"sealed":false'; then
        log_ok "bao-1 is unsealed and active"
        break
    fi
    if [[ "$i" -eq "$STARTUP_TIMEOUT" ]]; then
        log_error "Timed out waiting for bao-1 (${STARTUP_TIMEOUT}s)"
        exit 5
    fi
    sleep 1
done

# Run bao-init bootstrap
log_info "Running bao-init bootstrap ..."
$COMPOSE_CMD -f "$COMPOSE_FILE" up bao-init
# bao-init is a one-shot container — wait for it to exit
for i in $(seq 1 30); do
    init_status=$($RT inspect --format '{{.State.Status}}' bao-init 2>/dev/null || echo "missing")
    if [[ "$init_status" == "exited" ]]; then
        init_exit=$($RT inspect --format '{{.State.ExitCode}}' bao-init 2>/dev/null || echo "1")
        if [[ "$init_exit" != "0" ]]; then
            log_error "bao-init exited with code $init_exit"
            $RT logs bao-init 2>&1 | tail -20 || true
            exit 5
        fi
        log_ok "bao-init bootstrap completed successfully"
        break
    fi
    if [[ "$i" -eq 30 ]]; then
        log_error "bao-init did not complete within 30s"
        exit 5
    fi
    sleep 2
done

# Start pg-bao (reads AppRole creds from bao-init shared volume)
log_info "Starting pg-bao ..."
$COMPOSE_CMD -f "$COMPOSE_FILE" up -d pg-bao

# Wait for pg-bao to be ready
log_info "Waiting for pg-bao to be ready ..."
for i in $(seq 1 60); do
    if $RT exec pg-tde-bao pg_isready -U postgres > /dev/null 2>&1; then
        log_ok "pg-bao is ready"
        break
    fi
    if [[ "$i" -eq 60 ]]; then
        log_error "Timed out waiting for pg-bao"
        $RT logs pg-tde-bao 2>&1 | tail -30 || true
        exit 5
    fi
    sleep 2
done

START=$(timer_start)

# ---------------------------------------------------------------------------
# Test execution
# ---------------------------------------------------------------------------

# Test 1: OpenBao Raft leader healthy
bao_api_ok 1 "OpenBao bao-1 sealed=false" "$BAO_PORT_1" "/v1/sys/health"

# Test 2: AppRole credentials written by bao-init
log_info "Test 2: AppRole credentials exist ..."
if $RT exec pg-tde-bao test -f /tmp/bao-init/role_id 2>/dev/null; then
    pass "Test 2: AppRole credentials exist"
else
    fail "Test 2: AppRole credentials not found in shared volume"
fi

# Test 3: PostgreSQL extension loaded
run_bao_test 3 "Extension loaded on pg-bao" \
    "SELECT extname FROM pg_extension WHERE extname = 'pg_vault_tde'"

# Test 4: Vault DEK acquisition from real OpenBao
log_info "Test 4: DEK acquisition from OpenBao ..."
if $RT exec -u postgres pg-tde-bao psql -U postgres -At -c \
        "SELECT pg_vault_tde_vault_fetch_dek();" 2>&1 | grep -q "t\|true\|fetch"; then
    pass "Test 4: DEK acquisition from OpenBao"
else
    # DEK may have been auto-fetched on startup; verify generation > 0
    gen=$($RT exec -u postgres pg-tde-bao psql -U postgres -At \
              -c "SELECT pg_vault_tde_key_generation();" 2>/dev/null || echo "0")
    if [[ "$gen" -ge 0 ]]; then
        pass "Test 4: DEK generation=$gen (auto-fetched on startup)"
    else
        fail "Test 4: DEK acquisition from OpenBao"
    fi
fi

# Test 5: Encryption/decryption round-trip
log_info "Test 5: Encrypted round-trip ..."
result=$($RT exec -u postgres pg-tde-bao psql -U postgres -At -c "
    CREATE TABLE IF NOT EXISTS bao_test (id int, secret text) USING encrypted_heap;
    INSERT INTO bao_test VALUES (1, 'openbao-verified-data');
    SELECT secret FROM bao_test WHERE id = 1;
    DROP TABLE bao_test;
" 2>&1 || true)
if echo "$result" | grep -q "openbao-verified-data"; then
    pass "Test 5: Encryption/decryption round-trip"
else
    fail "Test 5: Encryption/decryption round-trip (got: $result)"
fi

# Test 6: health_check() shows vault_reachable=true
log_info "Test 6: health_check() vault_reachable=true ..."
health=$($RT exec -u postgres pg-tde-bao psql -U postgres -At -c \
    "SELECT vault_reachable FROM pg_vault_tde_health_check();" 2>/dev/null || echo "f")
if [[ "$health" == "t" ]]; then
    pass "Test 6: health_check vault_reachable=true"
else
    fail "Test 6: health_check vault_reachable returned: $health"
fi

# Test 7: health_check() shows auth_method=approle
log_info "Test 7: health_check() auth_method=approle ..."
auth_method=$($RT exec -u postgres pg-tde-bao psql -U postgres -At -c \
    "SELECT auth_method FROM pg_vault_tde_health_check();" 2>/dev/null || echo "unknown")
if [[ "$auth_method" == "approle" ]]; then
    pass "Test 7: health_check auth_method=approle"
else
    fail "Test 7: health_check auth_method returned: $auth_method"
fi

# Test 8: KEK rewrap — advances key version in OpenBao Raft
log_info "Test 8: pg_vault_tde_vault_rewrap_dek() ..."
if $RT exec -u postgres pg-tde-bao psql -U postgres -At -c \
        "SELECT pg_vault_tde_vault_rewrap_dek();" > /dev/null 2>&1; then
    pass "Test 8: KEK rewrap succeeded"
else
    fail "Test 8: KEK rewrap failed"
fi

# Test 9: BGW token renewal GUC is active
log_info "Test 9: BGW token renewal enabled ..."
bgw_enabled=$($RT exec -u postgres pg-tde-bao psql -U postgres -At -c \
    "SHOW pg_vault_tde.bgw_enabled;" 2>/dev/null || echo "off")
if [[ "$bgw_enabled" == "on" ]]; then
    pass "Test 9: BGW enabled (bgw_enabled=on)"
else
    fail "Test 9: BGW enabled check returned: $bgw_enabled"
fi

# Test 10: Multi-type data round-trip under Vault-provided DEK
log_info "Test 10: Multi-type encrypted table ..."
result=$($RT exec -u postgres pg-tde-bao psql -U postgres -At -c "
    CREATE TABLE bao_multi (
        id   serial PRIMARY KEY,
        name text,
        val  numeric(10,2),
        flag boolean,
        ts   timestamptz DEFAULT now()
    ) USING encrypted_heap;
    INSERT INTO bao_multi (name, val, flag)
        SELECT md5(g::text), g * 0.75, (g % 2 = 0)
        FROM generate_series(1, 50) g;
    SELECT count(*) FROM bao_multi;
    DROP TABLE bao_multi;
" 2>/dev/null || echo "0")
if [[ "$result" == "50" ]]; then
    pass "Test 10: Multi-type round-trip (50 rows)"
else
    fail "Test 10: Multi-type round-trip returned: $result"
fi

# Test 11: bao-2 follower is up and unsealed
bao_api_ok 11 "OpenBao bao-2 sealed=false (Raft follower)" "$BAO_PORT_2" "/v1/sys/health"

# Test 12: bao-3 follower is up and unsealed
bao_api_ok 12 "OpenBao bao-3 sealed=false (Raft follower)" "$BAO_PORT_3" "/v1/sys/health"

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
ELAPSED=$(timer_elapsed "$START")
TOTAL=$(( TESTS_PASSED + TESTS_FAILED ))
echo ""
log_stage "OPENBAO TEST RESULTS"
printf "  Passed: %d / %d   ($(timer_fmt "$ELAPSED"))\n" "$TESTS_PASSED" "$TOTAL"

if [[ $TESTS_FAILED -gt 0 ]]; then
    log_error "OPENBAO: $TESTS_FAILED test(s) FAILED"
    exit 5
fi

log_ok "OPENBAO: ALL $TESTS_PASSED TESTS PASSED ($(timer_fmt "$ELAPSED"))"
exit 0