#[pg_test]
fn set_payload_updates_jsonb_paths_and_typed_columns() {
    create_payload_mutation_collection("m12_payload_set");

    let rows = payload_mutation_rows(
        "SELECT source_key, updated
           FROM pgcontext.set_payload(
                'm12_payload_set',
                ARRAY['1'],
                '{\"topic\":\"postgres\",\"status\":\"closed\",\"priority\":5,\"archived\":true}'::jsonb
           )",
    );
    assert_eq!(rows, vec![("1".to_owned(), true)]);

    let row = payload_row("m12_payload_set", 1);
    assert_eq!(row.status.as_deref(), Some("closed"));
    assert_eq!(row.priority, Some(5));
    assert_eq!(row.archived, Some(true));
    assert_eq!(row.metadata, serde_json::json!({"topic":"postgres"}));
}

#[pg_test]
fn delete_payload_removes_jsonb_paths_and_nullable_columns() {
    create_payload_mutation_collection("m12_payload_delete");

    let rows = payload_mutation_rows(
        "SELECT source_key, updated
           FROM pgcontext.delete_payload(
                'm12_payload_delete',
                ARRAY['1'],
                ARRAY['topic', 'priority', 'archived']
           )",
    );
    assert_eq!(rows, vec![("1".to_owned(), true)]);

    let row = payload_row("m12_payload_delete", 1);
    assert_eq!(row.status.as_deref(), Some("open"));
    assert_eq!(row.priority, None);
    assert_eq!(row.archived, None);
    assert_eq!(row.metadata, serde_json::json!({}));
}

#[pg_test]
fn clear_payload_removes_registered_nullable_payloads() {
    create_payload_mutation_collection("m12_payload_clear");
    Spi::run(
        "SELECT pgcontext.delete_payload('m12_payload_clear', ARRAY['1'], ARRAY['status'])",
    )
    .expect("status should be nullable before clear");

    let rows = payload_mutation_rows(
        "SELECT source_key, updated
           FROM pgcontext.clear_payload('m12_payload_clear', ARRAY['1'])",
    );
    assert_eq!(rows, vec![("1".to_owned(), true)]);

    let row = payload_row("m12_payload_clear", 1);
    assert_eq!(row.status, None);
    assert_eq!(row.priority, None);
    assert_eq!(row.archived, None);
    assert_eq!(row.metadata, serde_json::json!({}));
}

#[pg_test]
#[should_panic(expected = "unknown payload field for collection m12_payload_unknown: missing")]
fn set_payload_rejects_unregistered_payload_keys() {
    create_payload_mutation_collection("m12_payload_unknown");

    Spi::run(
        "SELECT pgcontext.set_payload(
            'm12_payload_unknown',
            ARRAY['1'],
            '{\"missing\":\"value\"}'::jsonb
        )",
    )
    .expect("unknown payload field should fail");
}

#[pg_test]
fn set_payload_rejects_unknown_points_without_partial_updates() {
    create_payload_mutation_collection("m12_payload_missing_point");

    let result = std::panic::catch_unwind(|| {
        Spi::run(
            "SELECT pgcontext.set_payload(
                'm12_payload_missing_point',
                ARRAY['1', 'missing'],
                '{\"topic\":\"changed\"}'::jsonb
            )",
        )
        .expect("missing source key should fail");
    });
    assert!(result.is_err(), "missing point should panic");
    assert_eq!(
        payload_row("m12_payload_missing_point", 1).metadata,
        serde_json::json!({"topic":"rust"})
    );
}

#[pg_test]
#[should_panic(expected = "payload field is not nullable for collection m12_payload_not_null: status")]
fn delete_payload_rejects_non_nullable_ordinary_columns() {
    create_payload_mutation_collection("m12_payload_not_null");
    Spi::run("ALTER TABLE public.m12_payload_not_null ALTER COLUMN status SET NOT NULL")
        .expect("status should become not null");

    Spi::run(
        "SELECT pgcontext.delete_payload(
            'm12_payload_not_null',
            ARRAY['1'],
            ARRAY['status']
        )",
    )
    .expect("non-null ordinary payload delete should fail");
}

#[derive(Debug, PartialEq)]
struct PayloadRow {
    status: Option<String>,
    priority: Option<i32>,
    archived: Option<bool>,
    metadata: serde_json::Value,
}

fn create_payload_mutation_collection(collection_name: &str) {
    Spi::run(&format!(
        "CREATE TABLE public.{collection_name} (
             id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
             embedding vector,
             status text,
             priority integer,
             archived boolean,
             metadata jsonb NOT NULL DEFAULT '{{}}'::jsonb
         )"
    ))
    .expect("payload fixture table should be created");
    Spi::run(&format!(
        "INSERT INTO public.{collection_name}
            (id, embedding, status, priority, archived, metadata)
         VALUES
            (1, '[1,0]'::vector, 'open', 1, false, '{{\"topic\":\"rust\"}}'::jsonb)"
    ))
    .expect("payload fixture row should be inserted");
    Spi::run(&format!(
        "SELECT pgcontext.create_collection('{collection_name}', 'public.{collection_name}')"
    ))
    .expect("payload fixture collection should be created");
    Spi::run(&format!(
        "SELECT pgcontext.register_vector('{collection_name}', 'embedding', 'embedding', 2, 'l2')"
    ))
    .expect("payload fixture vector should be registered");
    Spi::run(&format!(
        "SELECT pgcontext.upsert_points('{collection_name}', ARRAY['1'])"
    ))
    .expect("payload fixture point should be registered");
    Spi::run(&format!(
        "SELECT pgcontext.register_filter_column('{collection_name}', 'status', 'status')"
    ))
    .expect("status payload should be registered");
    Spi::run(&format!(
        "SELECT pgcontext.register_filter_column('{collection_name}', 'priority', 'priority')"
    ))
    .expect("priority payload should be registered");
    Spi::run(&format!(
        "SELECT pgcontext.register_filter_column('{collection_name}', 'archived', 'archived')"
    ))
    .expect("archived payload should be registered");
    Spi::run(&format!(
        "SELECT pgcontext.register_jsonb_path('{collection_name}', 'topic', 'metadata', ARRAY['topic'])"
    ))
    .expect("topic payload should be registered");
}

fn payload_mutation_rows(sql: &str) -> Vec<(String, bool)> {
    Spi::connect(|client| {
        let rows = client.select(sql, None, &[])?;
        let mut output = Vec::new();
        for row in rows {
            output.push((
                row.get::<String>(1)?.expect("source_key should not be null"),
                row.get::<bool>(2)?.expect("updated should not be null"),
            ));
        }
        Ok::<_, spi::Error>(output)
    })
    .expect("payload mutation query should succeed")
}

fn payload_row(table_name: &str, id: i64) -> PayloadRow {
    Spi::connect(|client| {
        let rows = client.select(
            &format!(
                "SELECT status, priority, archived, metadata
                   FROM public.{table_name}
                  WHERE id = {id}"
            ),
            Some(1),
            &[],
        )?;
        let row = rows.first();
        Ok::<_, spi::Error>(PayloadRow {
            status: row.get::<String>(1)?,
            priority: row.get::<i32>(2)?,
            archived: row.get::<bool>(3)?,
            metadata: row
                .get::<pgrx::JsonB>(4)?
                .expect("metadata should not be null")
                .0,
        })
    })
    .expect("payload row query should succeed")
}
