//! SQL schema objects for pgContext-owned late-interaction token storage.

#[allow(
    unused_imports,
    reason = "pgrx resolves this import as an extension_sql dependency"
)]
use crate::vector::Vector;

pgrx::extension_sql!(
    r#"
CREATE TABLE pgcontext._collection_late_interaction (
    collection_id bigint PRIMARY KEY
        REFERENCES pgcontext._collections(collection_id) ON DELETE CASCADE,
    source_table_oid oid NOT NULL,
    source_schema_name text NOT NULL,
    source_table_name text NOT NULL,
    token_column_name text NOT NULL,
    token_attnum int2 NOT NULL,
    dimensions int4,
    hnsw_index_oid oid,
    point_count bigint NOT NULL DEFAULT 0 CHECK (point_count >= 0),
    token_count bigint NOT NULL DEFAULT 0 CHECK (token_count >= 0),
    status text NOT NULL DEFAULT 'building'
        CHECK (status IN ('building', 'ready', 'stale', 'failed')),
    created_at timestamptz NOT NULL DEFAULT pg_catalog.now(),
    updated_at timestamptz NOT NULL DEFAULT pg_catalog.now(),
    CHECK (dimensions IS NULL OR dimensions > 0),
    CHECK ((status = 'ready') = (dimensions IS NOT NULL AND hnsw_index_oid IS NOT NULL))
);

ALTER TABLE pgcontext._collection_points
    ADD CONSTRAINT pgcontext_collection_points_collection_point_unique
    UNIQUE (collection_id, point_id);

CREATE TABLE pgcontext._collection_late_interaction_tokens (
    token_id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
    collection_id bigint NOT NULL,
    point_id bigint NOT NULL,
    token_ordinal int4 NOT NULL CHECK (token_ordinal > 0),
    token_vector vector NOT NULL,
    created_at timestamptz NOT NULL DEFAULT pg_catalog.now(),
    updated_at timestamptz NOT NULL DEFAULT pg_catalog.now(),
    FOREIGN KEY (collection_id, point_id)
        REFERENCES pgcontext._collection_points(collection_id, point_id)
        ON DELETE CASCADE,
    UNIQUE (collection_id, point_id, token_ordinal)
);

CREATE VIEW pgcontext._visible_collection_late_interaction
WITH (security_barrier = true) AS
SELECT registrations.*
  FROM pgcontext._collection_late_interaction AS registrations
  JOIN pgcontext._collections AS collections USING (collection_id)
 WHERE pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER');

GRANT SELECT ON pgcontext._visible_collection_late_interaction TO PUBLIC;

CREATE FUNCTION pgcontext._capture_late_interaction_tokens()
RETURNS trigger
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    registration pgcontext._collection_late_interaction%ROWTYPE;
    selected_record record;
    current_source_key text;
    previous_source_key text;
    token_vectors pgcontext.vector[];
    point bigint;
    minimum_dimensions int4;
    maximum_dimensions int4;
    deleted_token_count bigint;
    replacement_token_count bigint;
BEGIN
    SELECT *
      INTO registration
      FROM pgcontext._collection_late_interaction
     WHERE collection_id = TG_ARGV[0]::bigint;

    IF NOT FOUND OR registration.source_table_oid <> TG_RELID THEN
        RAISE EXCEPTION 'late-interaction source trigger binding is stale for relation %', TG_RELID
            USING ERRCODE = '55000';
    END IF;

    IF TG_OP IN ('UPDATE', 'DELETE') THEN
        previous_source_key := pg_catalog.to_jsonb(OLD)->>'id';
        IF previous_source_key IS NULL OR previous_source_key = '' THEN
            RAISE EXCEPTION 'late-interaction source key id must not be null or empty'
                USING ERRCODE = '22023';
        END IF;
    END IF;

    IF TG_OP = 'DELETE' THEN
        DELETE FROM pgcontext._collection_late_interaction_tokens AS tokens
         USING pgcontext._collection_points AS points
         WHERE tokens.collection_id = registration.collection_id
           AND tokens.point_id = points.point_id
           AND points.collection_id = registration.collection_id
           AND points.source_key = previous_source_key;
        GET DIAGNOSTICS deleted_token_count = ROW_COUNT;

        IF deleted_token_count > 0 THEN
            UPDATE pgcontext._collection_late_interaction
               SET point_count = point_count - 1,
                   token_count = token_count - deleted_token_count,
                   updated_at = pg_catalog.now()
             WHERE collection_id = registration.collection_id;
        END IF;

        UPDATE pgcontext._collection_points
           SET deleted_at = coalesce(deleted_at, pg_catalog.now()),
               updated_at = pg_catalog.now()
         WHERE collection_id = registration.collection_id
           AND source_key = previous_source_key;
        RETURN OLD;
    END IF;

    current_source_key := pg_catalog.to_jsonb(NEW)->>'id';
    IF current_source_key IS NULL OR current_source_key = '' THEN
        RAISE EXCEPTION 'late-interaction source key id must not be null or empty'
            USING ERRCODE = '22023';
    END IF;

    EXECUTE pg_catalog.format(
        'SELECT ($1).%I::pgcontext.vector[] AS token_vectors',
        registration.token_column_name
    )
    INTO selected_record
    USING NEW;
    token_vectors := selected_record.token_vectors;

    IF token_vectors IS NULL
       OR pg_catalog.cardinality(token_vectors) = 0
       OR pg_catalog.array_position(token_vectors, NULL) IS NOT NULL THEN
        RAISE EXCEPTION 'late-interaction token source must contain at least one non-null vector for source key %', current_source_key
            USING ERRCODE = '22023';
    END IF;
    IF pg_catalog.cardinality(token_vectors) > 16384 THEN
        RAISE EXCEPTION 'late-interaction token count exceeds per-point limit 16384 for source key %', current_source_key
            USING ERRCODE = '54000';
    END IF;

    SELECT pg_catalog.min(pgcontext.vector_dims(token)),
           pg_catalog.max(pgcontext.vector_dims(token))
      INTO minimum_dimensions, maximum_dimensions
      FROM pg_catalog.unnest(token_vectors) AS token;
    IF minimum_dimensions IS NULL OR minimum_dimensions <> maximum_dimensions THEN
        RAISE EXCEPTION 'late-interaction token dimensions must be uniform for source key %', current_source_key
            USING ERRCODE = '22023';
    END IF;
    IF registration.dimensions IS NOT NULL
       AND registration.dimensions <> minimum_dimensions THEN
        RAISE EXCEPTION 'late-interaction token dimension mismatch: expected %, found % for source key %',
            registration.dimensions,
            minimum_dimensions,
            current_source_key
            USING ERRCODE = '22023';
    END IF;

    IF TG_OP = 'UPDATE' AND previous_source_key <> current_source_key THEN
        DELETE FROM pgcontext._collection_late_interaction_tokens AS tokens
         USING pgcontext._collection_points AS points
         WHERE tokens.collection_id = registration.collection_id
           AND tokens.point_id = points.point_id
           AND points.collection_id = registration.collection_id
           AND points.source_key = previous_source_key;
        GET DIAGNOSTICS deleted_token_count = ROW_COUNT;
        IF deleted_token_count > 0 THEN
            UPDATE pgcontext._collection_late_interaction
               SET point_count = point_count - 1,
                   token_count = token_count - deleted_token_count,
                   updated_at = pg_catalog.now()
             WHERE collection_id = registration.collection_id;
        END IF;
        UPDATE pgcontext._collection_points
           SET deleted_at = coalesce(deleted_at, pg_catalog.now()),
               updated_at = pg_catalog.now()
         WHERE collection_id = registration.collection_id
           AND source_key = previous_source_key;
    END IF;

    INSERT INTO pgcontext._collection_points (collection_id, source_key)
    VALUES (registration.collection_id, current_source_key)
    ON CONFLICT (collection_id, source_key) DO UPDATE
        SET deleted_at = NULL,
            updated_at = pg_catalog.now()
    RETURNING point_id INTO point;

    DELETE FROM pgcontext._collection_late_interaction_tokens
     WHERE collection_id = registration.collection_id
       AND point_id = point;
    GET DIAGNOSTICS replacement_token_count = ROW_COUNT;
    INSERT INTO pgcontext._collection_late_interaction_tokens (
        collection_id,
        point_id,
        token_ordinal,
        token_vector
    )
    SELECT registration.collection_id,
           point,
           ordinal::int4,
           token
      FROM pg_catalog.unnest(token_vectors) WITH ORDINALITY AS expanded(token, ordinal);

    UPDATE pgcontext._collection_late_interaction
       SET dimensions = coalesce(dimensions, minimum_dimensions),
           point_count = point_count + CASE WHEN replacement_token_count = 0 THEN 1 ELSE 0 END,
           token_count = token_count
               - replacement_token_count
               + pg_catalog.cardinality(token_vectors),
           updated_at = pg_catalog.now()
     WHERE collection_id = registration.collection_id;
    RETURN NEW;
END;
$$;

CREATE FUNCTION pgcontext._begin_late_interaction_registration(
    p_collection_id bigint,
    p_source_table_oid oid,
    p_source_schema_name text,
    p_source_table_name text,
    p_token_column_name text,
    p_token_attnum int2
)
RETURNS void
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    trigger_name text;
BEGIN
    IF NOT EXISTS (
        SELECT 1
          FROM pgcontext._collections AS collections
          JOIN pg_catalog.pg_class AS source_class
            ON source_class.oid = collections.source_table_oid
          JOIN pg_catalog.pg_namespace AS source_namespace
            ON source_namespace.oid = source_class.relnamespace
          JOIN pg_catalog.pg_attribute AS token_attribute
            ON token_attribute.attrelid = source_class.oid
           AND token_attribute.attname = p_token_column_name
           AND token_attribute.attnum = p_token_attnum
           AND token_attribute.attnum > 0
           AND NOT token_attribute.attisdropped
         WHERE collections.collection_id = p_collection_id
           AND pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER')
           AND source_class.oid = p_source_table_oid
           AND source_namespace.nspname = p_source_schema_name
           AND source_class.relname = p_source_table_name
           AND source_class.relkind = 'r'
           AND token_attribute.atttypid = 'pgcontext.vector[]'::pg_catalog.regtype
    ) THEN
        RAISE EXCEPTION 'invalid or unauthorized late-interaction registration for collection %', p_collection_id
            USING ERRCODE = '42501';
    END IF;
    IF NOT pg_catalog.has_table_privilege(SESSION_USER, p_source_table_oid, 'SELECT') THEN
        RAISE EXCEPTION 'permission denied for late-interaction source table %.%',
            p_source_schema_name,
            p_source_table_name
            USING ERRCODE = '42501';
    END IF;

    INSERT INTO pgcontext._collection_late_interaction (
        collection_id,
        source_table_oid,
        source_schema_name,
        source_table_name,
        token_column_name,
        token_attnum
    )
    VALUES (
        p_collection_id,
        p_source_table_oid,
        p_source_schema_name,
        p_source_table_name,
        p_token_column_name,
        p_token_attnum
    );

    trigger_name := pg_catalog.format('pgcontext_late_interaction_%s', p_collection_id);
    EXECUTE pg_catalog.format(
        'CREATE TRIGGER %I AFTER INSERT OR UPDATE OF id, %I OR DELETE ON %I.%I '
        'FOR EACH ROW EXECUTE FUNCTION pgcontext._capture_late_interaction_tokens(%L)',
        trigger_name,
        p_token_column_name,
        p_source_schema_name,
        p_source_table_name,
        p_collection_id::text
    );
END;
$$;

CREATE FUNCTION pgcontext._store_late_interaction_tokens(
    p_collection_id bigint,
    p_source_key text
)
RETURNS bigint
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    registration pgcontext._collection_late_interaction%ROWTYPE;
    point bigint;
    minimum_dimensions int4;
    maximum_dimensions int4;
    replacement_token_count bigint;
    token_vectors pgcontext.vector[];
    loaded_row_count bigint;
BEGIN
    SELECT registrations.*
      INTO registration
      FROM pgcontext._collection_late_interaction AS registrations
      JOIN pgcontext._collections AS collections USING (collection_id)
     WHERE registrations.collection_id = p_collection_id
       AND pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER');
    IF NOT FOUND
       OR NOT pg_catalog.has_table_privilege(
           SESSION_USER,
           registration.source_table_oid,
           'SELECT'
       ) THEN
        RAISE EXCEPTION 'permission denied for late-interaction collection %', p_collection_id
            USING ERRCODE = '42501';
    END IF;
    IF p_source_key IS NULL OR p_source_key = '' THEN
        RAISE EXCEPTION 'late-interaction source key must not be null or empty'
            USING ERRCODE = '22023';
    END IF;

    EXECUTE pg_catalog.format(
        'SELECT source.%I::pgcontext.vector[]
           FROM %I.%I AS source
          WHERE source.id::text = $1
          LIMIT 1',
        registration.token_column_name,
        registration.source_schema_name,
        registration.source_table_name
    ) INTO token_vectors USING p_source_key;
    GET DIAGNOSTICS loaded_row_count = ROW_COUNT;
    IF loaded_row_count = 0 THEN
        RAISE EXCEPTION 'late-interaction source row does not exist for source key %', p_source_key
            USING ERRCODE = '42704';
    END IF;
    IF token_vectors IS NULL
       OR pg_catalog.cardinality(token_vectors) = 0
       OR pg_catalog.array_position(token_vectors, NULL) IS NOT NULL THEN
        RAISE EXCEPTION 'late-interaction token source must contain at least one non-null vector for source key %', p_source_key
            USING ERRCODE = '22023';
    END IF;
    IF pg_catalog.cardinality(token_vectors) > 16384 THEN
        RAISE EXCEPTION 'late-interaction token count exceeds per-point limit 16384 for source key %', p_source_key
            USING ERRCODE = '54000';
    END IF;

    SELECT pg_catalog.min(pgcontext.vector_dims(token)),
           pg_catalog.max(pgcontext.vector_dims(token))
      INTO minimum_dimensions, maximum_dimensions
      FROM pg_catalog.unnest(token_vectors) AS token;
    IF minimum_dimensions IS NULL OR minimum_dimensions <> maximum_dimensions THEN
        RAISE EXCEPTION 'late-interaction token dimensions must be uniform for source key %', p_source_key
            USING ERRCODE = '22023';
    END IF;
    IF registration.dimensions IS NOT NULL
       AND registration.dimensions <> minimum_dimensions THEN
        RAISE EXCEPTION 'late-interaction token dimension mismatch: expected %, found % for source key %',
            registration.dimensions,
            minimum_dimensions,
            p_source_key
            USING ERRCODE = '22023';
    END IF;

    INSERT INTO pgcontext._collection_points (collection_id, source_key)
    VALUES (p_collection_id, p_source_key)
    ON CONFLICT (collection_id, source_key) DO UPDATE
        SET updated_at = pg_catalog.now()
    RETURNING point_id INTO point;

    DELETE FROM pgcontext._collection_late_interaction_tokens
     WHERE collection_id = p_collection_id
       AND point_id = point;
    GET DIAGNOSTICS replacement_token_count = ROW_COUNT;
    INSERT INTO pgcontext._collection_late_interaction_tokens (
        collection_id,
        point_id,
        token_ordinal,
        token_vector
    )
    SELECT p_collection_id,
           point,
           ordinal::int4,
           token
      FROM pg_catalog.unnest(token_vectors) WITH ORDINALITY AS expanded(token, ordinal);

    UPDATE pgcontext._collection_late_interaction
       SET dimensions = coalesce(dimensions, minimum_dimensions),
           point_count = point_count + CASE WHEN replacement_token_count = 0 THEN 1 ELSE 0 END,
           token_count = token_count
               - replacement_token_count
               + pg_catalog.cardinality(token_vectors),
           updated_at = pg_catalog.now()
     WHERE collection_id = p_collection_id;
    RETURN point;
END;
$$;

CREATE FUNCTION pgcontext._finish_late_interaction_registration(
    p_collection_id bigint,
    p_dimensions int4
)
RETURNS oid
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    index_name text;
    index_oid oid;
BEGIN
    IF p_dimensions <= 0 OR NOT EXISTS (
        SELECT 1
          FROM pgcontext._collection_late_interaction AS registrations
          JOIN pgcontext._collections AS collections USING (collection_id)
         WHERE registrations.collection_id = p_collection_id
           AND registrations.dimensions = p_dimensions
           AND pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER')
    ) THEN
        RAISE EXCEPTION 'invalid or unauthorized late-interaction finalization for collection %', p_collection_id
            USING ERRCODE = '42501';
    END IF;

    index_name := pg_catalog.format('pgcontext_late_interaction_%s_hnsw', p_collection_id);
    EXECUTE pg_catalog.format(
        'CREATE INDEX %I ON pgcontext._collection_late_interaction_tokens '
        'USING pgcontext_hnsw ((token_vector::pgcontext.vector(%s)) pgcontext.vector_hnsw_ip_ops) '
        'WHERE collection_id = %s',
        index_name,
        p_dimensions,
        p_collection_id
    );
    index_oid := pg_catalog.to_regclass(
        pg_catalog.format('pgcontext.%I', index_name)
    )::oid;
    UPDATE pgcontext._collection_late_interaction
       SET hnsw_index_oid = index_oid,
           status = 'ready',
           updated_at = pg_catalog.now()
     WHERE collection_id = p_collection_id;
    RETURN index_oid;
END;
$$;

CREATE FUNCTION pgcontext._late_interaction_ann_candidate_points(
    p_collection_id bigint,
    p_query vector,
    p_limit int4
)
RETURNS TABLE (point_id bigint)
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    registration pgcontext._collection_late_interaction%ROWTYPE;
    previous_enable_seqscan text;
BEGIN
    IF p_limit IS NULL OR p_limit <= 0 THEN
        RAISE EXCEPTION 'late-interaction ANN candidate limit must be positive'
            USING ERRCODE = '22023';
    END IF;

    SELECT registrations.*
      INTO registration
      FROM pgcontext._collection_late_interaction AS registrations
      JOIN pgcontext._collections AS collections USING (collection_id)
      JOIN pg_catalog.pg_index AS index_metadata
        ON index_metadata.indexrelid = registrations.hnsw_index_oid
      JOIN pg_catalog.pg_class AS index_class
        ON index_class.oid = index_metadata.indexrelid
      JOIN pg_catalog.pg_am AS access_method
        ON access_method.oid = index_class.relam
      JOIN pg_catalog.pg_opclass AS operator_class
        ON operator_class.oid = index_metadata.indclass[0]
      JOIN pg_catalog.pg_namespace AS operator_namespace
        ON operator_namespace.oid = operator_class.opcnamespace
     WHERE registrations.collection_id = p_collection_id
       AND pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER')
       AND registrations.status = 'ready'
       AND registrations.dimensions IS NOT NULL
       AND index_metadata.indrelid = 'pgcontext._collection_late_interaction_tokens'::regclass
       AND index_metadata.indisvalid
       AND index_metadata.indisready
       AND access_method.amname = 'pgcontext_hnsw'
       AND index_class.relname = pg_catalog.format(
           'pgcontext_late_interaction_%s_hnsw',
           registrations.collection_id
       )
       AND index_metadata.indnkeyatts = 1
       AND index_metadata.indnatts = 1
       AND operator_namespace.nspname = 'pgcontext'
       AND operator_class.opcname = 'vector_hnsw_ip_ops'
       AND pg_catalog.regexp_replace(
           pg_catalog.pg_get_expr(
               index_metadata.indpred,
               index_metadata.indrelid,
               true
           ),
           '[()[:space:]]',
           '',
           'g'
       ) = pg_catalog.format('collection_id=%s', registrations.collection_id)
       AND pg_catalog.regexp_replace(
           pg_catalog.pg_get_indexdef(index_metadata.indexrelid, 1, true),
           '[()[:space:]]',
           '',
           'g'
       ) IN (
           pg_catalog.format('token_vector::vector%s', registrations.dimensions),
           pg_catalog.format('token_vector::pgcontext.vector%s', registrations.dimensions)
       );
    IF NOT FOUND THEN
        RAISE EXCEPTION 'late-interaction ANN generation is not ready or is unauthorized for collection %', p_collection_id
            USING ERRCODE = '55000';
    END IF;
    IF pgcontext.vector_dims(p_query) <> registration.dimensions THEN
        RAISE EXCEPTION 'late-interaction query dimension mismatch: expected %, found %',
            registration.dimensions,
            pgcontext.vector_dims(p_query)
            USING ERRCODE = '22023';
    END IF;

    previous_enable_seqscan := pg_catalog.current_setting('enable_seqscan');
    PERFORM pg_catalog.set_config('enable_seqscan', 'off', true);
    BEGIN
        RETURN QUERY EXECUTE pg_catalog.format(
            'SELECT tokens.point_id
               FROM pgcontext._collection_late_interaction_tokens AS tokens
              WHERE tokens.collection_id = $1
              ORDER BY (tokens.token_vector::pgcontext.vector(%s)) OPERATOR(pgcontext.<#>) $2
              LIMIT $3',
            registration.dimensions
        ) USING p_collection_id, p_query, p_limit;
    EXCEPTION WHEN others THEN
        PERFORM pg_catalog.set_config('enable_seqscan', previous_enable_seqscan, true);
        RAISE;
    END;
    PERFORM pg_catalog.set_config('enable_seqscan', previous_enable_seqscan, true);
END;
$$;

CREATE FUNCTION pgcontext._prepare_late_interaction_repair(
    p_collection_id bigint,
    p_source_table_oid oid,
    p_token_attnum int2
)
RETURNS void
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    registration pgcontext._collection_late_interaction%ROWTYPE;
    trigger_name text;
    index_name text;
BEGIN
    SELECT registrations.*
      INTO registration
      FROM pgcontext._collection_late_interaction AS registrations
      JOIN pgcontext._collections AS collections USING (collection_id)
      JOIN pg_catalog.pg_class AS source_class
        ON source_class.oid = p_source_table_oid
      JOIN pg_catalog.pg_namespace AS source_namespace
        ON source_namespace.oid = source_class.relnamespace
      JOIN pg_catalog.pg_attribute AS token_attribute
        ON token_attribute.attrelid = source_class.oid
       AND token_attribute.attname = registrations.token_column_name
       AND token_attribute.attnum = p_token_attnum
       AND token_attribute.attnum > 0
       AND NOT token_attribute.attisdropped
     WHERE registrations.collection_id = p_collection_id
       AND collections.source_table_oid = p_source_table_oid
       AND pg_catalog.pg_has_role(SESSION_USER, collections.owner_role, 'MEMBER')
       AND source_namespace.nspname = registrations.source_schema_name
       AND source_class.relname = registrations.source_table_name
       AND source_class.relkind = 'r'
       AND token_attribute.atttypid = 'pgcontext.vector[]'::pg_catalog.regtype;
    IF NOT FOUND OR NOT pg_catalog.has_table_privilege(
        SESSION_USER,
        p_source_table_oid,
        'SELECT'
    ) THEN
        RAISE EXCEPTION 'invalid or unauthorized late-interaction repair for collection %', p_collection_id
            USING ERRCODE = '42501';
    END IF;

    trigger_name := pg_catalog.format('pgcontext_late_interaction_%s', p_collection_id);
    EXECUTE pg_catalog.format(
        'DROP TRIGGER IF EXISTS %I ON %I.%I',
        trigger_name,
        registration.source_schema_name,
        registration.source_table_name
    );
    index_name := pg_catalog.format('pgcontext_late_interaction_%s_hnsw', p_collection_id);
    EXECUTE pg_catalog.format('DROP INDEX IF EXISTS pgcontext.%I', index_name);

    DELETE FROM pgcontext._collection_late_interaction_tokens
     WHERE collection_id = p_collection_id;
    UPDATE pgcontext._collection_late_interaction
       SET source_table_oid = p_source_table_oid,
           token_attnum = p_token_attnum,
           dimensions = NULL,
           hnsw_index_oid = NULL,
           point_count = 0,
           token_count = 0,
           status = 'building',
           updated_at = pg_catalog.now()
     WHERE collection_id = p_collection_id;

    EXECUTE pg_catalog.format(
        'CREATE TRIGGER %I AFTER INSERT OR UPDATE OF id, %I OR DELETE ON %I.%I '
        'FOR EACH ROW EXECUTE FUNCTION pgcontext._capture_late_interaction_tokens(%L)',
        trigger_name,
        registration.token_column_name,
        registration.source_schema_name,
        registration.source_table_name,
        p_collection_id::text
    );
END;
$$;

CREATE FUNCTION pgcontext._cleanup_late_interaction_registration()
RETURNS trigger
LANGUAGE plpgsql
SECURITY DEFINER
SET search_path = pg_catalog, pgcontext
AS $$
DECLARE
    registration pgcontext._collection_late_interaction%ROWTYPE;
    trigger_name text;
    index_name text;
    current_schema_name text;
    current_table_name text;
BEGIN
    SELECT *
      INTO registration
      FROM pgcontext._collection_late_interaction
     WHERE collection_id = OLD.collection_id;
    IF NOT FOUND THEN
        RETURN OLD;
    END IF;

    trigger_name := pg_catalog.format(
        'pgcontext_late_interaction_%s',
        OLD.collection_id
    );
    SELECT namespace.nspname, source_class.relname
      INTO current_schema_name, current_table_name
      FROM pg_catalog.pg_class AS source_class
      JOIN pg_catalog.pg_namespace AS namespace
        ON namespace.oid = source_class.relnamespace
     WHERE source_class.oid = registration.source_table_oid;
    IF FOUND THEN
        EXECUTE pg_catalog.format(
            'DROP TRIGGER IF EXISTS %I ON %I.%I',
            trigger_name,
            current_schema_name,
            current_table_name
        );
    END IF;

    index_name := pg_catalog.format(
        'pgcontext_late_interaction_%s_hnsw',
        OLD.collection_id
    );
    EXECUTE pg_catalog.format('DROP INDEX IF EXISTS pgcontext.%I', index_name);
    RETURN OLD;
END;
$$;

CREATE TRIGGER pgcontext_cleanup_late_interaction_registration
BEFORE DELETE ON pgcontext._collections
FOR EACH ROW
EXECUTE FUNCTION pgcontext._cleanup_late_interaction_registration();

SELECT pg_catalog.pg_extension_config_dump(
    'pgcontext._collection_late_interaction',
    ''
);
SELECT pg_catalog.pg_extension_config_dump(
    'pgcontext._collection_late_interaction_tokens',
    ''
);
"#,
    name = "create_late_interaction_catalog_tables",
    requires = [Vector, "create_catalog_tables"]
);
