#!/usr/bin/env bash
# The findings of the external audit of 0.2.4 that this repo closed, each against its
# control -- the proof that the instrument can answer the other way.
#
#   RG-01 the sampled vector was cast back to the column's type, so on a domain column the
#         cast ran the domain's CHECK -- a function its owner wrote -- as whoever measured.
#   RG-02 measure() read the parent table without ONLY: inheritance children, which the
#         parent's index does not cover, diluted the measurement of a broken index.
#   RG-03 the plan check looked for "Index Scan", not for the index asked about: with a twin
#         index on the same column, measure(X) reported the twin's recall.
#   RG-04 rows were told apart by ctid, which repeats across partitions.
#   RG-08 baselines and measurements were not dumped: after a restore check() returned no
#         rows, which reads as "all fine".
#
# Needs pgvector. Run against the throwaway cluster: test/cluster.sh init && start.

set -euo pipefail

PG_CONFIG=${PG_CONFIG:-pg_config}
BIN=$("$PG_CONFIG" --bindir)
PSQL=${PSQL:-$BIN/psql}
RAIZ=$(cd "$(dirname "$0")/.." && pwd)
export PGHOST=${PGHOST:-$RAIZ/.testcluster} PGPORT=${PGPORT:-5496}
DB=recall_guard_test_audit
RESTORED=recall_guard_test_audit_restored
TENANT=recall_guard_test_audit_tenant
DUMP=$RAIZ/.testcluster/audit.dump
failures=0

if [ ! -f "$("$PG_CONFIG" --sharedir)/extension/vector.control" ]; then
    echo "DID NOT RUN: pgvector is not installed in this PostgreSQL" >&2
    exit 2
fi
for d in "$DB" "$RESTORED"; do
    if [ "$($PSQL -X -d postgres -tAc "select 1 from pg_database where datname = '$d'")" = 1 ]; then
        echo "a database $d already exists: not dropping it, somebody else made it" >&2
        exit 2
    fi
done
if [ "$($PSQL -X -d postgres -tAc "select 1 from pg_roles where rolname = '$TENANT'")" = 1 ]; then
    echo "a role $TENANT already exists: not dropping it, somebody else made it" >&2
    exit 2
fi
cleanup() {
    $PSQL -X -d postgres -qc "drop database if exists $DB" -c "drop database if exists $RESTORED" \
        -c "drop role if exists $TENANT" >/dev/null 2>&1 || true
    rm -f "$DUMP"
}
trap cleanup EXIT

q() { $PSQL -X -d "$DB" -tA "$@" 2>&1 || true; }
qt() { PGUSER=$TENANT $PSQL -X -d "$DB" -tA "$@" 2>&1 || true; }

check() {
    local what="$1" expected="$2" got="$3"
    if [[ "$got" == *"$expected"* ]]; then
        echo "  ok   $what"
    else
        echo "  FAIL $what"
        echo "       expected: $expected"
        echo "       got:      $got"
        failures=$((failures + 1))
    fi
}

$PSQL -X -d postgres -qc "create database $DB" -c "create role $TENANT login"
$PSQL -X -d "$DB" -q -v ON_ERROR_STOP=1 -v tenant="$TENANT" >/dev/null <<'SQL'
CREATE EXTENSION vector;
CREATE EXTENSION pg_recall_guard CASCADE;
CREATE SCHEMA tenant_s AUTHORIZATION :"tenant";
SELECT setseed(0.42);
SQL

echo "RG-01: measuring a domain column runs none of its owner's code"
qt -q -v ON_ERROR_STOP=1 >/dev/null <<'SQL'
CREATE TABLE tenant_s.ran_as (who text);
CREATE FUNCTION tenant_s.chk(v point) RETURNS boolean LANGUAGE plpgsql AS $f$
BEGIN
    INSERT INTO tenant_s.ran_as VALUES (current_user);
    RETURN true;
END $f$;
CREATE DOMAIN tenant_s.pt AS point;
CREATE TABLE tenant_s.notes (id serial, e tenant_s.pt);
INSERT INTO tenant_s.notes (e) SELECT point(random() * 100, random() * 100) FROM generate_series(1, 500);
CREATE INDEX notes_gist ON tenant_s.notes USING gist (e);
GRANT USAGE ON SCHEMA tenant_s TO PUBLIC;
GRANT SELECT, INSERT ON tenant_s.ran_as TO PUBLIC;
GRANT SELECT ON tenant_s.notes TO PUBLIC;
SQL
q -q -c "analyze tenant_s.notes" -c "select recall_guard.approve('tenant_s.notes_gist', 10, 20, 'reviewed')" >/dev/null
qt -q -c "alter domain tenant_s.pt add constraint c check (tenant_s.chk(value))" -c "truncate tenant_s.ran_as" >/dev/null
check "control: casting to the domain as the superuser runs the owner's function" "as_superuser=1" \
    "$(q -c "select '(1,1)'::tenant_s.pt" >/dev/null; q -c "select 'as_superuser=' || count(*) from tenant_s.ran_as where who <> '$TENANT'")"
q -c "delete from tenant_s.ran_as" >/dev/null
check "check() by the superuser runs none of it" "as_superuser=0" \
    "$(q -c "select verdict from recall_guard.check() where index_name = 'tenant_s.notes_gist'" >/dev/null; q -c "select 'as_superuser=' || count(*) from tenant_s.ran_as where who <> '$TENANT'")"
check "  ...and still measures the index" "ok" \
    "$(q -c "select verdict from recall_guard.check() where index_name = 'tenant_s.notes_gist'")"

echo "RG-02: inheritance children the index does not cover do not dilute the measurement"
q -q -v ON_ERROR_STOP=1 >/dev/null <<'SQL'
CREATE TABLE inh_parent (id int, emb vector(16));
CREATE TABLE inh_child () INHERITS (inh_parent);
INSERT INTO inh_parent SELECT g, (SELECT array_agg(random())::vector(16) FROM generate_series(1, 16) WHERE g > 0) FROM generate_series(1, 300) g;
INSERT INTO inh_child SELECT g, (SELECT array_agg(random())::vector(16) FROM generate_series(1, 16) WHERE g > 0) FROM generate_series(301, 6300) g;
CREATE INDEX inh_hnsw ON inh_parent USING hnsw (emb vector_l2_ops);
CREATE TABLE flat (id int, emb vector(16));
INSERT INTO flat SELECT id, emb FROM ONLY inh_parent;
CREATE INDEX flat_hnsw ON flat USING hnsw (emb vector_l2_ops);
ANALYZE;
SQL
flat=$(q -c "set hnsw.ef_search = 1" -c "select round(avg(recall_guard.measure('flat_hnsw', 10, 300)), 2) from generate_series(1, 3)" | tail -1)
check "control: the same rows in a plain table, ef_search = 1, measure low" "low=true" \
    "$(q -c "select 'low=' || ($flat < 0.5)")"
check "the parent's index measures the same, children or not" "low=true" \
    "$(q -c "set hnsw.ef_search = 1" -c "select 'low=' || (round(avg(recall_guard.measure('inh_hnsw', 10, 300)), 2) < 0.5) from generate_series(1, 3)" | tail -1)"

echo "RG-03: measure(X) measures X, not the index the planner prefers"
q -q -v ON_ERROR_STOP=1 >/dev/null <<'SQL'
CREATE TABLE twin (id int, emb vector(16));
INSERT INTO twin SELECT g, (SELECT array_agg(random())::vector(16) FROM generate_series(1, 16) WHERE g > 0) FROM generate_series(1, 5000) g;
CREATE INDEX twin_ivf ON twin USING ivfflat (emb vector_l2_ops) WITH (lists = 70);
CREATE INDEX twin_hnsw ON twin USING hnsw (emb vector_l2_ops);
ANALYZE twin;
SQL
plan=$(q -c "set hnsw.ef_search = 1" -c "set ivfflat.probes = 1" -c "set enable_seqscan = off" -c "explain (costs off) select id from twin order by emb <-> (select emb from twin limit 1) limit 10" | grep -o "using twin_[a-z]*" | head -1)
check "control: with probes = 1 the planner prefers the ivfflat twin" "twin_ivf" "$plan"
check "measure(twin_hnsw) does not report the twin's recall: it says which index was read" "did not measure public.twin_hnsw: read twin_ivf" \
    "$(q -c "set hnsw.ef_search = 1" -c "set ivfflat.probes = 1" -c "select recall_guard.measure('twin_hnsw', 10, 50)" | grep -o 'no midió el índice [^:]*: leyó [a-z_]*' | sed 's/no midió el índice/did not measure/; s/leyó/read/')"

echo "RG-04: rows of different partitions are different rows"
q -q -v ON_ERROR_STOP=1 >/dev/null <<'SQL'
CREATE TABLE parts (id int, emb vector(4)) PARTITION BY HASH (id);
DO $$ BEGIN FOR i IN 0..49 LOOP
    EXECUTE format('CREATE TABLE parts_%s PARTITION OF parts FOR VALUES WITH (MODULUS 50, REMAINDER %s)', i, i);
END LOOP; END $$;
INSERT INTO parts SELECT g, (SELECT array_agg(random())::vector(4) FROM generate_series(1, 4) WHERE g > 0) FROM generate_series(1, 1000) g;
CREATE INDEX parts_ivf ON parts USING ivfflat (emb vector_l2_ops) WITH (lists = 2);
ANALYZE parts;
-- The recall by id, the way it should come out: per row, the k nearest other rows through the
-- index against the exact k nearest, matched by id.
CREATE TEMP TABLE by_index AS SELECT 1;
SQL
truth=$(q -q -c "set ivfflat.probes = 1" -c "set enable_seqscan = off" \
    -c "create table idx_ids as select p.id, array(select q.id from parts q where q.id <> p.id order by q.emb <-> p.emb limit 10) ids from parts p" \
    -c "reset enable_seqscan" -c "set enable_indexscan = off" -c "set enable_bitmapscan = off" \
    -c "create table exact_ids as select p.id, array(select q.id from parts q where q.id <> p.id order by q.emb <-> p.emb limit 10) ids from parts p" \
    -c "select round(avg((select count(*) from unnest(i.ids) x where x = any(e.ids))::numeric / 10), 2) from idx_ids i join exact_ids e using (id)" | tail -1)
check "control: the true recall by id is below 1" "below=true" "$(q -c "select 'below=' || ($truth < 1)")"
check "measure() over the whole table is the recall by id" "measured=$truth" \
    "$(q -c "set ivfflat.probes = 1" -c "select 'measured=' || round(recall_guard.measure('parts_ivf', 10, 1000), 2)" | tail -1)"

echo "RG-08: baselines and measurements survive pg_dump and restore"
q -q -c "select recall_guard.approve('flat_hnsw', 10, 30)" >/dev/null
"$BIN/pg_dump" -Fc -d "$DB" -f "$DUMP"
$PSQL -X -d postgres -qc "create database $RESTORED"
"$BIN/pg_restore" -d "$RESTORED" "$DUMP" >/dev/null 2>&1 || true
check "control: the source has baselines" "baselines=true" \
    "$(q -c "select 'baselines=' || (count(*) > 0) from recall_guard.baselines")"
check "the restored database has the same baselines" "same=true" \
    "$($PSQL -X -d "$RESTORED" -tAc "select 'same=' || (count(*) = $(q -c "select count(*) from recall_guard.baselines")) from recall_guard.baselines" 2>&1)"
check "  ...and its history" "history=true" \
    "$($PSQL -X -d "$RESTORED" -tAc "select 'history=' || (count(*) > 0) from recall_guard.measurements" 2>&1)"
check "  ...and a new measurement there still gets a new id" "ok" \
    "$($PSQL -X -d "$RESTORED" -tAc "select recall_guard.measure('flat_hnsw', 10, 10) is not null" -c "select 'ok'" 2>&1 | tail -1)"

if [ "$failures" -ne 0 ]; then
    echo "$failures check(s) failed"
    exit 1
fi
echo "the findings of the 0.2.4 audit are closed, each against its control"