#[pg_test]
fn rolled_back_create_collection_leaves_no_catalog_row() {
    Spi::run(
        r#"
        DO $$
        BEGIN
            BEGIN
                PERFORM pgcontext.create_collection('m2_tx_create_rollback');
                RAISE EXCEPTION 'force catalog rollback';
            EXCEPTION WHEN others THEN
                NULL;
            END;
        END $$;
        "#,
    )
    .expect("rollback block should complete");

    assert_eq!(rollback_collection_count("m2_tx_create_rollback"), 0);
}

#[pg_test]
fn rolled_back_register_vector_leaves_no_vector_row() {
    rollback_create_source_table("m2_tx_vector_rollback", "embedding vector");
    Spi::run(
        "SELECT pgcontext.create_collection('m2_tx_vector_rollback', 'public.m2_tx_vector_rollback')",
    )
    .expect("table-backed collection should be created");

    Spi::run(
        r#"
        DO $$
        BEGIN
            BEGIN
                PERFORM pgcontext.register_vector(
                    'm2_tx_vector_rollback',
                    'embedding',
                    'embedding',
                    3,
                    'l2'
                );
                RAISE EXCEPTION 'force vector rollback';
            EXCEPTION WHEN others THEN
                NULL;
            END;
        END $$;
        "#,
    )
    .expect("rollback block should complete");

    assert_eq!(rollback_vector_count("m2_tx_vector_rollback"), 0);
}

#[pg_test]
fn rolled_back_upsert_points_leaves_no_point_rows() {
    rollback_create_point_collection("m2_tx_point_upsert_rollback");

    Spi::run(
        r#"
        DO $$
        BEGIN
            BEGIN
                PERFORM pgcontext.upsert_points(
                    'm2_tx_point_upsert_rollback',
                    ARRAY['doc-1', 'doc-2']
                );
                RAISE EXCEPTION 'force point rollback';
            EXCEPTION WHEN others THEN
                NULL;
            END;
        END $$;
        "#,
    )
    .expect("rollback block should complete");

    assert_eq!(rollback_point_count("m2_tx_point_upsert_rollback"), 0);
}

#[pg_test]
fn rolled_back_delete_points_preserves_active_point() {
    rollback_create_point_collection("m2_tx_point_delete_rollback");
    Spi::run(
        "SELECT pgcontext.upsert_points('m2_tx_point_delete_rollback', ARRAY['doc-1'])",
    )
    .expect("point should be inserted");

    Spi::run(
        r#"
        DO $$
        BEGIN
            BEGIN
                PERFORM pgcontext.delete_points(
                    'm2_tx_point_delete_rollback',
                    ARRAY['doc-1']
                );
                RAISE EXCEPTION 'force point delete rollback';
            EXCEPTION WHEN others THEN
                NULL;
            END;
        END $$;
        "#,
    )
    .expect("rollback block should complete");

    assert_eq!(
        rollback_active_point_count("m2_tx_point_delete_rollback", "doc-1"),
        1
    );
}

fn rollback_create_point_collection(collection_name: &str) {
    rollback_create_source_table(collection_name, "embedding vector");
    Spi::run(&format!(
        "SELECT pgcontext.create_collection('{collection_name}', 'public.{collection_name}')"
    ))
    .expect("point collection should be created");
}

fn rollback_create_source_table(table_name: &str, vector_column_sql: &str) {
    Spi::run(&format!(
        "CREATE TABLE public.{table_name} (
             id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
             {vector_column_sql}
         )"
    ))
    .expect("source table should be created");
}

fn rollback_collection_count(collection_name: &str) -> i64 {
    rollback_count_query(&format!(
        "SELECT count(*)::bigint
           FROM pgcontext._collections
          WHERE collection_name = '{collection_name}'"
    ))
}

fn rollback_vector_count(collection_name: &str) -> i64 {
    rollback_count_query(&format!(
        "SELECT count(*)::bigint
           FROM pgcontext._collection_vectors vectors
           JOIN pgcontext._collections collections USING (collection_id)
          WHERE collections.collection_name = '{collection_name}'"
    ))
}

fn rollback_point_count(collection_name: &str) -> i64 {
    rollback_count_query(&format!(
        "SELECT count(*)::bigint
           FROM pgcontext._collection_points points
           JOIN pgcontext._collections collections USING (collection_id)
          WHERE collections.collection_name = '{collection_name}'"
    ))
}

fn rollback_active_point_count(collection_name: &str, source_key: &str) -> i64 {
    rollback_count_query(&format!(
        "SELECT count(*)::bigint
           FROM pgcontext._collection_points points
           JOIN pgcontext._collections collections USING (collection_id)
          WHERE collections.collection_name = '{collection_name}'
            AND points.source_key = '{source_key}'
            AND points.deleted_at IS NULL"
    ))
}

fn rollback_count_query(sql: &str) -> i64 {
    Spi::get_one::<i64>(sql)
        .expect("count query should succeed")
        .expect("count should not be null")
}
