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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DBNAME="${DBNAME:-pgcontext_late_interaction_ann}"
# shellcheck source=tests/heavy/lib.sh
source "${SCRIPT_DIR}/lib.sh"

start_and_install_extension
reset_database

psql_db <<'SQL'
CREATE EXTENSION pgcontext;

CREATE TABLE public.late_ann_docs (
    id bigint PRIMARY KEY,
    token_vectors vector[] NOT NULL
);

INSERT INTO public.late_ann_docs (id, token_vectors)
VALUES (10, ARRAY['[1,0]'::vector, '[0,1]'::vector]),
       (20, ARRAY['[0.8,0.1]'::vector, '[0.1,0.7]'::vector]),
       (30, ARRAY['[1,0]'::vector, '[1,0]'::vector]),
       (40, ARRAY['[0,0]'::vector]);

SELECT pgcontext.create_collection('late_ann_docs', 'public.late_ann_docs');
SELECT pgcontext.upsert_points('late_ann_docs', ARRAY['10', '20', '30', '40']);

CREATE TABLE public.late_ann_tokens (
    source_key text NOT NULL,
    token_embedding vector(2) NOT NULL
);

INSERT INTO public.late_ann_tokens (source_key, token_embedding)
VALUES ('10', '[1,0]'::vector),
       ('10', '[0,1]'::vector),
       ('20', '[1,0]'::vector),
       ('20', '[0,1]'::vector),
       ('30', '[1,0]'::vector),
       ('30', '[0,1]'::vector),
       ('40', '[1,0]'::vector);

CREATE INDEX late_ann_tokens_embedding_idx
    ON public.late_ann_tokens USING pgcontext_hnsw (token_embedding);

DO $$
DECLARE
    actual_keys text[];
    actual_scores numeric[];
    deleted_count integer;
BEGIN
    WITH ranked AS (
        SELECT row_number() OVER () AS ord, source_key, score
          FROM pgcontext.search_late_interaction_ann(
              'late_ann_docs',
              ARRAY['[1,0]'::vector, '[0,1]'::vector],
              'token_vectors',
              'public.late_ann_tokens',
              'source_key',
              'token_embedding',
              3,
              3
          )
    )
    SELECT array_agg(source_key ORDER BY ord),
           array_agg(round(score::numeric, 3) ORDER BY ord)
      INTO actual_keys, actual_scores
      FROM ranked;

    IF actual_keys <> ARRAY['10', '20', '30'] THEN
        RAISE EXCEPTION 'late-interaction ANN candidate keys mismatch: %',
            actual_keys;
    END IF;
    IF actual_scores <> ARRAY[2.000, 1.500, 1.000]::numeric[] THEN
        RAISE EXCEPTION 'late-interaction ANN rerank scores mismatch: %',
            actual_scores;
    END IF;

    UPDATE public.late_ann_docs
       SET token_vectors = ARRAY['[0,0]'::vector]
     WHERE id = 20;

    WITH ranked AS (
        SELECT row_number() OVER () AS ord, source_key, score
          FROM pgcontext.search_late_interaction_ann(
              'late_ann_docs',
              ARRAY['[1,0]'::vector, '[0,1]'::vector],
              'token_vectors',
              'public.late_ann_tokens',
              'source_key',
              'token_embedding',
              10,
              3
          )
    )
    SELECT array_agg(source_key ORDER BY ord),
           array_agg(round(score::numeric, 3) ORDER BY ord)
      INTO actual_keys, actual_scores
      FROM ranked;

    IF actual_keys <> ARRAY['10', '30', '20'] THEN
        RAISE EXCEPTION 'late-interaction ANN source recheck keys mismatch: %',
            actual_keys;
    END IF;
    IF actual_scores <> ARRAY[2.000, 1.000, 0.000]::numeric[] THEN
        RAISE EXCEPTION 'late-interaction ANN source recheck scores mismatch: %',
            actual_scores;
    END IF;

    PERFORM pgcontext.delete_points('late_ann_docs', ARRAY['30']);

    WITH ranked AS (
        SELECT source_key
          FROM pgcontext.search_late_interaction_ann(
              'late_ann_docs',
              ARRAY['[1,0]'::vector, '[0,1]'::vector],
              'token_vectors',
              'public.late_ann_tokens',
              'source_key',
              'token_embedding',
              10,
              4
          )
    )
    SELECT count(*) FILTER (WHERE source_key = '30'),
           array_agg(source_key ORDER BY source_key)
      INTO deleted_count, actual_keys
      FROM ranked;

    IF deleted_count <> 0 THEN
        RAISE EXCEPTION 'late-interaction ANN returned deleted source key 30';
    END IF;
    IF actual_keys <> ARRAY['10', '20', '40'] THEN
        RAISE EXCEPTION 'late-interaction ANN deleted recheck keys mismatch: %',
            actual_keys;
    END IF;
END
$$;

CREATE TABLE public.late_ann_budget_docs (
    id bigint PRIMARY KEY,
    token_vectors vector[] NOT NULL
);

INSERT INTO public.late_ann_budget_docs (id, token_vectors)
VALUES (10, array_fill('[1,0]'::vector, ARRAY[1000]));

SELECT pgcontext.create_collection(
    'late_ann_budget_docs',
    'public.late_ann_budget_docs'
);
SELECT pgcontext.upsert_points('late_ann_budget_docs', ARRAY['10']);

CREATE TABLE public.late_ann_budget_tokens (
    source_key text NOT NULL,
    token_embedding vector(2) NOT NULL
);

INSERT INTO public.late_ann_budget_tokens (source_key, token_embedding)
VALUES ('10', '[1,0]'::vector);

CREATE INDEX late_ann_budget_tokens_embedding_idx
    ON public.late_ann_budget_tokens USING pgcontext_hnsw (token_embedding);

DO $$
BEGIN
    PERFORM *
      FROM pgcontext.search_late_interaction_ann(
          'late_ann_budget_docs',
          array_fill('[1,0]'::vector, ARRAY[1001]),
          'token_vectors',
          'public.late_ann_budget_tokens',
          'source_key',
          'token_embedding',
          1000,
          1
      );
    RAISE EXCEPTION 'late-interaction ANN budget rejection did not fire';
EXCEPTION
    WHEN SQLSTATE '54000' THEN
        IF SQLERRM <> 'late interaction comparison budget exceeded: 1001000 > 1000000' THEN
            RAISE EXCEPTION 'unexpected late-interaction ANN budget message: %',
                SQLERRM;
        END IF;
END
$$;
SQL

printf 'late_interaction_ann_candidates_deduped\n'
printf 'late_interaction_ann_exact_rerank_scores\n'
printf 'late_interaction_ann_source_recheck\n'
printf 'late_interaction_ann_deleted_recheck\n'
printf 'late_interaction_ann_budget_rejected\n'