use std::collections::HashMap;
use std::sync::Arc;

use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray, UInt64Array};
use arrow_schema::{DataType, Field, Schema};
use async_trait::async_trait;
use lance::dataset::builder::DatasetBuilder;
use lance_namespace::LanceNamespace;
use lance_namespace::models::{
    DescribeTableRequest, DescribeTableVersionRequest, ListTableVersionsRequest,
};
use lance_namespace_impls::DirectoryNamespaceBuilder;
use tokio::sync::Mutex;

use super::publisher::ManifestBatchPublisher;
use super::*;
use omnigraph_compiler::catalog::build_catalog;
use omnigraph_compiler::schema::parser::parse_schema;

fn test_schema_source() -> &'static str {
    r#"
node Person {
    name: String
    age: I32?
}
node Company {
    name: String
}
edge Knows: Person -> Person {
    since: Date?
}
edge WorksAt: Person -> Company {
    title: String?
}
"#
}

fn build_test_catalog() -> Catalog {
    let schema = parse_schema(test_schema_source()).unwrap();
    build_catalog(&schema).unwrap()
}

#[tokio::test]
async fn test_init_creates_manifest_and_sub_tables() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();

    assert!(snap.entry("node:Person").is_some());
    assert!(snap.entry("node:Company").is_some());
    assert!(snap.entry("edge:Knows").is_some());
    assert!(snap.entry("edge:WorksAt").is_some());

    for key in &["node:Person", "node:Company", "edge:Knows", "edge:WorksAt"] {
        let entry = snap.entry(key).unwrap();
        assert_eq!(entry.table_version, 1);
        assert_eq!(entry.row_count, 0);
        assert!(entry.table_branch.is_none());
    }
}

#[tokio::test]
async fn test_open_reads_existing_manifest() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    ManifestCoordinator::init(uri, &catalog).await.unwrap();

    let mc = ManifestCoordinator::open(uri).await.unwrap();
    let snap = mc.snapshot();
    assert!(snap.entry("node:Person").is_some());
    assert!(snap.entry("edge:Knows").is_some());
}

#[tokio::test]
async fn test_commit_advances_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let v1 = mc.version();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;

    let new_version = mc
        .commit(&[SubTableUpdate {
            table_key: "node:Person".to_string(),
            table_version: person_version,
            table_branch: None,
            row_count: 1,
            version_metadata: table_version_metadata_for_state(
                uri,
                &person_entry.table_path,
                None,
                person_version,
            )
            .await
            .unwrap(),
        }])
        .await
        .unwrap();

    assert!(new_version > v1);

    let snap = mc.snapshot();
    let person = snap.entry("node:Person").unwrap();
    assert_eq!(person.table_version, person_version);
    assert_eq!(person.row_count, 1);

    let company = snap.entry("node:Company").unwrap();
    assert_eq!(company.table_version, 1);
    assert_eq!(company.row_count, 0);
}

#[tokio::test]
async fn test_commit_changes_can_register_new_table_and_tombstone_old_one() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let before_version = mc.version();
    let person_entry = mc.snapshot().entry("node:Person").unwrap().clone();

    let table_key = "node:Human".to_string();
    let table_path = table_path_for_table_key(&table_key).unwrap();
    let dataset_uri = format!("{}/{}", uri, table_path);
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Utf8, false),
        Field::new("name", DataType::Utf8, false),
        Field::new("age", DataType::Int32, true),
    ]));
    let ds = crate::table_store::TableStore::create_empty_dataset(&dataset_uri, &schema)
        .await
        .unwrap();
    let state = crate::table_store::TableStore::new(uri)
        .table_state(&dataset_uri, &ds)
        .await
        .unwrap();

    mc.commit_changes(&[
        ManifestChange::RegisterTable(TableRegistration {
            table_key: table_key.clone(),
            table_path: table_path.clone(),
        }),
        ManifestChange::Update(SubTableUpdate {
            table_key: table_key.clone(),
            table_version: state.version,
            table_branch: None,
            row_count: state.row_count,
            version_metadata: state.version_metadata,
        }),
        ManifestChange::Tombstone(TableTombstone {
            table_key: "node:Person".to_string(),
            tombstone_version: person_entry.table_version + 1,
        }),
    ])
    .await
    .unwrap();

    let head = mc.snapshot();
    assert!(head.entry("node:Human").is_some());
    assert!(head.entry("node:Person").is_none());

    let historical = ManifestCoordinator::snapshot_at(uri, None, before_version)
        .await
        .unwrap();
    assert!(historical.entry("node:Person").is_some());
    assert!(historical.entry("node:Human").is_none());
}

#[tokio::test]
async fn test_snapshot_open_sub_table() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_ds = snap.open("node:Person").await.unwrap();

    assert_eq!(person_ds.schema().fields.len(), 3);
    assert_eq!(person_ds.count_rows(None).await.unwrap(), 0);
}

#[tokio::test]
async fn test_version_is_manifest_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    assert_eq!(mc.version(), snap.version());
}

#[tokio::test]
async fn test_list_branches_only_returns_main_once() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let branches = mc.list_branches().await.unwrap();
    assert_eq!(
        branches
            .iter()
            .filter(|branch| branch.as_str() == "main")
            .count(),
        1
    );
}

#[tokio::test]
async fn test_branch_namespace_lists_and_describes_versions() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();

    let namespace = branch_manifest_namespace(uri, None);
    let request =
        version_metadata.to_create_table_version_request("node:Person", person_version, 1, None);
    namespace.create_table_version(request).await.unwrap();
    mc.refresh().await.unwrap();

    let versions = namespace
        .list_table_versions(ListTableVersionsRequest {
            id: Some(vec!["node:Person".to_string()]),
            descending: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    assert_eq!(versions.versions.len(), 2);
    assert_eq!(versions.versions[0].version as u64, person_version);
    assert_eq!(versions.versions[1].version, 1);

    let described = namespace
        .describe_table_version(DescribeTableVersionRequest {
            id: Some(vec!["node:Person".to_string()]),
            version: Some(person_version as i64),
            ..Default::default()
        })
        .await
        .unwrap();
    assert_eq!(described.version.version as u64, person_version);
    assert_eq!(
        mc.snapshot().entry("node:Person").unwrap().table_version,
        person_version
    );
    assert_eq!(mc.snapshot().entry("node:Person").unwrap().row_count, 1);
}

#[tokio::test]
async fn test_directory_namespace_direct_publish_cannot_replace_native_omnigraph_write_path() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();

    let namespace = DirectoryNamespaceBuilder::new(uri)
        .manifest_enabled(true)
        .dir_listing_enabled(false)
        .table_version_tracking_enabled(true)
        .table_version_storage_enabled(true)
        .inline_optimization_enabled(false)
        .build()
        .await
        .unwrap();

    let versions = namespace
        .list_table_versions(ListTableVersionsRequest {
            id: Some(vec!["node:Person".to_string()]),
            descending: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    assert_eq!(
        versions.versions[0].version as u64,
        person_entry.table_version
    );

    let err = namespace
        .describe_table_version(DescribeTableVersionRequest {
            id: Some(vec!["node:Person".to_string()]),
            version: Some(person_version as i64),
            ..Default::default()
        })
        .await
        .unwrap_err();
    assert!(err.to_string().contains("not found"));

    let err = namespace
        .create_table_version(version_metadata.to_create_table_version_request(
            "node:Person",
            person_version,
            1,
            None,
        ))
        .await
        .unwrap_err();
    assert!(err.to_string().contains("already exists"));

    mc.refresh().await.unwrap();
    assert_eq!(
        mc.snapshot().entry("node:Person").unwrap().table_version,
        person_entry.table_version
    );
    assert_eq!(mc.snapshot().entry("node:Person").unwrap().row_count, 0);
}

#[tokio::test]
async fn test_snapshot_at_reads_branch_pinned_historical_state() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let main_manifest_version = mc.version();
    mc.create_branch("feature").await.unwrap();

    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    person_ds
        .create_branch("feature", person_entry.table_version, None)
        .await
        .unwrap();
    let mut feature_ds = person_ds.checkout_branch("feature").await.unwrap();
    let person_schema = Arc::new(feature_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    feature_ds.append(reader, None).await.unwrap();
    let feature_version = feature_ds.version().version;
    let feature_metadata = table_version_metadata_for_state(
        uri,
        &person_entry.table_path,
        Some("feature"),
        feature_version,
    )
    .await
    .unwrap();

    let namespace = branch_manifest_namespace(uri, Some("feature"));
    let request = feature_metadata.to_create_table_version_request(
        "node:Person",
        feature_version,
        1,
        Some("feature"),
    );
    namespace.create_table_version(request).await.unwrap();

    let feature_mc = ManifestCoordinator::open_at_branch(uri, "feature")
        .await
        .unwrap();
    let feature_snapshot =
        ManifestCoordinator::snapshot_at(uri, Some("feature"), feature_mc.version())
            .await
            .unwrap();
    let feature_entry = feature_snapshot.entry("node:Person").unwrap();
    assert_eq!(feature_entry.table_version, feature_version);
    assert_eq!(feature_entry.table_branch.as_deref(), Some("feature"));
    assert_eq!(
        feature_snapshot
            .open("node:Person")
            .await
            .unwrap()
            .count_rows(None)
            .await
            .unwrap(),
        1
    );

    let main_snapshot = ManifestCoordinator::snapshot_at(uri, None, main_manifest_version)
        .await
        .unwrap();
    let main_entry = main_snapshot.entry("node:Person").unwrap();
    assert_eq!(main_entry.table_version, person_entry.table_version);
    assert_eq!(main_entry.table_branch, None);
    assert_eq!(
        main_snapshot
            .open("node:Person")
            .await
            .unwrap()
            .count_rows(None)
            .await
            .unwrap(),
        0
    );
}

#[tokio::test]
async fn test_branch_manifest_namespace_uses_entry_owner_branch_for_latest_table_reads() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    mc.create_branch("feature").await.unwrap();

    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let company_entry = snap.entry("node:Company").unwrap().clone();

    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    person_ds
        .create_branch("feature", person_entry.table_version, None)
        .await
        .unwrap();
    let mut feature_person_ds = person_ds.checkout_branch("feature").await.unwrap();
    let person_schema = Arc::new(feature_person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    feature_person_ds.append(reader, None).await.unwrap();
    let feature_person_version = feature_person_ds.version().version;
    let feature_person_metadata = table_version_metadata_for_state(
        uri,
        &person_entry.table_path,
        Some("feature"),
        feature_person_version,
    )
    .await
    .unwrap();

    branch_manifest_namespace(uri, Some("feature"))
        .create_table_version(feature_person_metadata.to_create_table_version_request(
            "node:Person",
            feature_person_version,
            1,
            Some("feature"),
        ))
        .await
        .unwrap();

    let feature_namespace = branch_manifest_namespace(uri, Some("feature"));

    let inherited_company = feature_namespace
        .describe_table(DescribeTableRequest {
            id: Some(vec!["node:Company".to_string()]),
            with_table_uri: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    let inherited_company_uri = inherited_company.table_uri.as_deref().unwrap();
    assert!(
        !inherited_company_uri.contains("/tree/feature"),
        "inherited table should resolve to its owning branch, got {inherited_company_uri}"
    );

    let branch_owned_person = feature_namespace
        .describe_table(DescribeTableRequest {
            id: Some(vec!["node:Person".to_string()]),
            with_table_uri: Some(true),
            ..Default::default()
        })
        .await
        .unwrap();
    let branch_owned_person_uri = branch_owned_person.table_uri.as_deref().unwrap();
    assert!(
        branch_owned_person_uri.contains("/tree/feature"),
        "branch-owned table should resolve to feature branch, got {branch_owned_person_uri}"
    );

    let inherited_company_ds = DatasetBuilder::from_namespace(
        Arc::clone(&feature_namespace),
        vec!["node:Company".to_string()],
    )
    .await
    .unwrap()
    .with_branch("feature", None)
    .load()
    .await
    .unwrap();
    assert_eq!(inherited_company_ds.count_rows(None).await.unwrap(), 0);

    let branch_owned_person_ds = DatasetBuilder::from_namespace(
        Arc::clone(&feature_namespace),
        vec!["node:Person".to_string()],
    )
    .await
    .unwrap()
    .with_branch("feature", None)
    .load()
    .await
    .unwrap();
    assert_eq!(branch_owned_person_ds.count_rows(None).await.unwrap(), 1);
    assert_eq!(
        company_entry.table_branch, None,
        "sanity check: company table stays inherited on feature"
    );
}

#[tokio::test]
async fn test_refresh_observes_external_publish_without_mutating_existing_snapshot() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut reader = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let frozen_snapshot = reader.snapshot();
    let person_entry = frozen_snapshot.entry("node:Person").unwrap().clone();
    let manifest_version = reader.version();

    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader_batch = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader_batch, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();

    branch_manifest_namespace(uri, None)
        .create_table_version(version_metadata.to_create_table_version_request(
            "node:Person",
            person_version,
            1,
            None,
        ))
        .await
        .unwrap();

    assert_eq!(reader.version(), manifest_version);
    assert_eq!(
        frozen_snapshot.entry("node:Person").unwrap().table_version,
        person_entry.table_version
    );
    assert_eq!(
        frozen_snapshot
            .open("node:Person")
            .await
            .unwrap()
            .count_rows(None)
            .await
            .unwrap(),
        0
    );

    reader.refresh().await.unwrap();
    assert!(reader.version() > manifest_version);
    assert_eq!(
        reader
            .snapshot()
            .entry("node:Person")
            .unwrap()
            .table_version,
        person_version
    );
    assert_eq!(reader.snapshot().entry("node:Person").unwrap().row_count, 1);
}

#[tokio::test]
async fn test_batch_create_table_versions_is_atomic_on_conflict() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let manifest_version = mc.version();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let company_entry = snap.entry("node:Company").unwrap().clone();

    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;

    let person_version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();
    let company_version_metadata = table_version_metadata_for_state(
        uri,
        &company_entry.table_path,
        None,
        company_entry.table_version,
    )
    .await
    .unwrap();

    let person_request = person_version_metadata.to_create_table_version_request(
        "node:Person",
        person_version,
        1,
        None,
    );

    let conflicting_company_request = company_version_metadata.to_create_table_version_request(
        "node:Company",
        company_entry.table_version,
        0,
        None,
    );

    let err = GraphNamespacePublisher::new(uri, None)
        .publish_requests(&[person_request, conflicting_company_request])
        .await
        .unwrap_err();
    assert!(err.to_string().contains("already exists"));

    let reopened = ManifestCoordinator::open(uri).await.unwrap();
    assert_eq!(reopened.version(), manifest_version);
    assert_eq!(
        reopened
            .snapshot()
            .entry("node:Person")
            .unwrap()
            .table_version,
        person_entry.table_version
    );
    assert_eq!(
        reopened.snapshot().entry("node:Person").unwrap().row_count,
        0
    );
}

#[tokio::test]
async fn test_batch_create_table_versions_rejects_duplicate_requests_without_advancing_manifest() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let manifest_version = mc.version();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();

    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();
    let request =
        version_metadata.to_create_table_version_request("node:Person", person_version, 1, None);

    let err = GraphNamespacePublisher::new(uri, None)
        .publish_requests(&[request.clone(), request])
        .await
        .unwrap_err();
    assert!(err.to_string().contains("already exists"));

    let reopened = ManifestCoordinator::open(uri).await.unwrap();
    assert_eq!(reopened.version(), manifest_version);
    assert_eq!(
        reopened
            .snapshot()
            .entry("node:Person")
            .unwrap()
            .table_version,
        person_entry.table_version
    );
    assert_eq!(
        reopened.snapshot().entry("node:Person").unwrap().row_count,
        0
    );
}

#[tokio::test]
async fn test_batch_create_table_versions_allows_owner_branch_handoff_at_same_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut main_mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    main_mc.create_branch("feature").await.unwrap();

    let snap = main_mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    person_ds
        .create_branch("feature", person_entry.table_version, None)
        .await
        .unwrap();
    let mut feature_ds = person_ds.checkout_branch("feature").await.unwrap();
    let person_schema = Arc::new(feature_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    feature_ds.append(reader, None).await.unwrap();
    let feature_version = feature_ds.version().version;
    let feature_metadata = table_version_metadata_for_state(
        uri,
        &person_entry.table_path,
        Some("feature"),
        feature_version,
    )
    .await
    .unwrap();

    branch_manifest_namespace(uri, Some("feature"))
        .create_table_version(feature_metadata.to_create_table_version_request(
            "node:Person",
            feature_version,
            1,
            Some("feature"),
        ))
        .await
        .unwrap();

    let mut feature_mc = ManifestCoordinator::open_at_branch(uri, "feature")
        .await
        .unwrap();
    feature_mc.create_branch("experiment").await.unwrap();
    feature_ds
        .create_branch("experiment", feature_version, None)
        .await
        .unwrap();
    let experiment_metadata = table_version_metadata_for_state(
        uri,
        &person_entry.table_path,
        Some("experiment"),
        feature_version,
    )
    .await
    .unwrap();

    GraphNamespacePublisher::new(uri, Some("experiment"))
        .publish_requests(&[experiment_metadata.to_create_table_version_request(
            "node:Person",
            feature_version,
            1,
            Some("experiment"),
        )])
        .await
        .unwrap();

    let experiment_mc = ManifestCoordinator::open_at_branch(uri, "experiment")
        .await
        .unwrap();
    let experiment_snapshot = experiment_mc.snapshot();
    let experiment_entry = experiment_snapshot.entry("node:Person").unwrap();
    assert_eq!(experiment_entry.table_version, feature_version);
    assert_eq!(experiment_entry.table_branch.as_deref(), Some("experiment"));
}

#[tokio::test]
async fn test_staged_namespace_lists_native_table_versions_before_publish() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();

    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;

    let namespace = staged_table_namespace(uri, "node:Person", &person_entry.table_path, None);
    let listed = namespace
        .list_table_versions(ListTableVersionsRequest {
            id: Some(vec!["node:Person".to_string()]),
            descending: Some(false),
            ..Default::default()
        })
        .await
        .unwrap();
    let listed_versions: Vec<u64> = listed
        .versions
        .into_iter()
        .map(|version| version.version as u64)
        .collect();
    assert_eq!(listed_versions, vec![1, person_version]);

    let described = namespace
        .describe_table_version(DescribeTableVersionRequest {
            id: Some(vec!["node:Person".to_string()]),
            version: Some(person_version as i64),
            ..Default::default()
        })
        .await
        .unwrap();
    assert_eq!(described.version.version as u64, person_version);
}

#[derive(Clone)]
struct RecordingPublisher {
    inner: Arc<GraphNamespacePublisher>,
    requests: Arc<Mutex<Vec<CreateTableVersionRequest>>>,
}

impl RecordingPublisher {
    fn new(root_uri: &str, branch: Option<&str>) -> Self {
        Self {
            inner: Arc::new(GraphNamespacePublisher::new(root_uri, branch)),
            requests: Arc::new(Mutex::new(Vec::new())),
        }
    }

    async fn recorded_requests(&self) -> Vec<CreateTableVersionRequest> {
        self.requests.lock().await.clone()
    }
}

#[async_trait]
impl ManifestBatchPublisher for RecordingPublisher {
    async fn publish(
        &self,
        changes: &[ManifestChange],
        expected_table_versions: &HashMap<String, u64>,
    ) -> Result<Dataset> {
        let requests: Vec<CreateTableVersionRequest> = changes
            .iter()
            .filter_map(|change| match change {
                ManifestChange::Update(update) => Some(update.to_create_table_version_request()),
                ManifestChange::RegisterTable(_) | ManifestChange::Tombstone(_) => None,
            })
            .collect();
        self.requests.lock().await.extend_from_slice(&requests);
        self.inner.publish(changes, expected_table_versions).await
    }
}

struct FailingPublisher;

#[async_trait]
impl ManifestBatchPublisher for FailingPublisher {
    async fn publish(
        &self,
        _changes: &[ManifestChange],
        _expected_table_versions: &HashMap<String, u64>,
    ) -> Result<Dataset> {
        Err(OmniError::manifest(
            "injected batch publisher failure".to_string(),
        ))
    }
}

#[tokio::test]
async fn test_commit_routes_through_injected_batch_publisher() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();

    let recording = RecordingPublisher::new(uri, None);
    mc = mc.with_batch_publisher(Arc::new(recording.clone()));

    mc.commit(&[SubTableUpdate {
        table_key: "node:Person".to_string(),
        table_version: person_version,
        table_branch: None,
        row_count: 1,
        version_metadata: version_metadata.clone(),
    }])
    .await
    .unwrap();

    let recorded = recording.recorded_requests().await;
    assert_eq!(recorded.len(), 1);
    let request = &recorded[0];
    assert_eq!(
        request.id.as_ref().unwrap(),
        &vec!["node:Person".to_string()]
    );
    assert_eq!(request.version as u64, person_version);
    assert_eq!(request.manifest_path, version_metadata.manifest_path());
    assert_eq!(
        request.manifest_size,
        version_metadata.manifest_size().map(|size| size as i64)
    );
    assert_eq!(request.e_tag.as_deref(), version_metadata.e_tag());
    assert_eq!(
        request.naming_scheme.as_deref(),
        version_metadata.naming_scheme()
    );
    assert_eq!(
        request
            .metadata
            .as_ref()
            .and_then(|metadata| metadata.get(OMNIGRAPH_ROW_COUNT_KEY))
            .map(String::as_str),
        Some("1")
    );
    assert_eq!(
        mc.snapshot().entry("node:Person").unwrap().table_version,
        person_version
    );
}

#[tokio::test]
async fn test_commit_failure_from_injected_batch_publisher_preserves_visible_state() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();

    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let manifest_version = mc.version();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let person_batch = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec!["person-1"])),
            Arc::new(StringArray::from(vec!["Alice"])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(person_batch)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let person_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, person_version)
            .await
            .unwrap();

    mc = mc.with_batch_publisher(Arc::new(FailingPublisher));
    let err = mc
        .commit(&[SubTableUpdate {
            table_key: "node:Person".to_string(),
            table_version: person_version,
            table_branch: None,
            row_count: 1,
            version_metadata,
        }])
        .await
        .unwrap_err();
    assert!(err.to_string().contains("injected batch publisher failure"));
    assert_eq!(mc.version(), manifest_version);
    assert_eq!(
        mc.snapshot().entry("node:Person").unwrap().table_version,
        person_entry.table_version
    );
    assert_eq!(mc.snapshot().entry("node:Person").unwrap().row_count, 0);

    let reopened = ManifestCoordinator::open(uri).await.unwrap();
    assert_eq!(reopened.version(), manifest_version);
    assert_eq!(
        reopened
            .snapshot()
            .entry("node:Person")
            .unwrap()
            .table_version,
        person_entry.table_version
    );
}

/// Drive Person to a fresh on-disk dataset version `v` (returns the new
/// version number) and produce a `SubTableUpdate` ready to publish.
async fn append_person_and_make_update(
    uri: &str,
    person_entry: &SubTableEntry,
    name: &str,
) -> SubTableUpdate {
    let mut person_ds = Dataset::open(&format!("{}/{}", uri, person_entry.table_path))
        .await
        .unwrap();
    let person_schema = Arc::new(person_ds.schema().into());
    let row = RecordBatch::try_new(
        Arc::clone(&person_schema),
        vec![
            Arc::new(StringArray::from(vec![format!("person-{name}")])),
            Arc::new(StringArray::from(vec![Some(name.to_string())])),
            Arc::new(Int32Array::from(vec![Some(30)])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(row)], person_schema);
    person_ds.append(reader, None).await.unwrap();
    let new_version = person_ds.version().version;
    let version_metadata =
        table_version_metadata_for_state(uri, &person_entry.table_path, None, new_version)
            .await
            .unwrap();
    SubTableUpdate {
        table_key: "node:Person".to_string(),
        table_version: new_version,
        table_branch: None,
        row_count: 1,
        version_metadata,
    }
}

#[tokio::test]
async fn test_commit_with_expected_accepts_matching_versions() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();

    let update = append_person_and_make_update(uri, &person_entry, "Alice").await;
    let mut expected = HashMap::new();
    // After init, every table is at table_version=1 — assert that.
    expected.insert("node:Person".to_string(), 1);
    expected.insert("node:Company".to_string(), 1);

    mc.commit_with_expected(&[update.clone()], &expected)
        .await
        .expect("matching expected versions should publish cleanly");

    let after = mc.snapshot();
    assert_eq!(
        after.entry("node:Person").unwrap().table_version,
        update.table_version
    );
}

#[tokio::test]
async fn test_commit_with_expected_rejects_stale_with_typed_details() {
    use crate::error::ManifestConflictDetails;

    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();

    // Writer A advances Person.
    let update_a = append_person_and_make_update(uri, &person_entry, "Alice").await;
    let advanced_version = update_a.table_version;
    mc.commit(&[update_a]).await.unwrap();

    // Writer B then tries to commit, asserting Person is still at v=1.
    let update_b = append_person_and_make_update(uri, &person_entry, "Bob").await;
    let mut stale_expected = HashMap::new();
    stale_expected.insert("node:Person".to_string(), 1);

    let err = mc
        .commit_with_expected(&[update_b], &stale_expected)
        .await
        .expect_err("stale expected_table_versions should reject");

    match err {
        OmniError::Manifest(m) => match m.details {
            Some(ManifestConflictDetails::ExpectedVersionMismatch {
                table_key,
                expected,
                actual,
            }) => {
                assert_eq!(table_key, "node:Person");
                assert_eq!(expected, 1);
                assert_eq!(actual, advanced_version);
            }
            other => panic!("expected ExpectedVersionMismatch details, got {:?}", other),
        },
        other => panic!("expected OmniError::Manifest, got {:?}", other),
    }
}

#[tokio::test]
async fn test_commit_with_expected_catches_drift_on_untouched_table() {
    use crate::error::ManifestConflictDetails;

    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let snap = mc.snapshot();
    let person_entry = snap.entry("node:Person").unwrap().clone();
    let company_entry = snap.entry("node:Company").unwrap().clone();

    // Writer A advances Company.
    let mut company_ds = Dataset::open(&format!("{}/{}", uri, company_entry.table_path))
        .await
        .unwrap();
    let company_schema = Arc::new(company_ds.schema().into());
    let row = RecordBatch::try_new(
        Arc::clone(&company_schema),
        vec![
            Arc::new(StringArray::from(vec!["company-1"])),
            Arc::new(StringArray::from(vec!["Acme"])),
        ],
    )
    .unwrap();
    let reader = RecordBatchIterator::new(vec![Ok(row)], company_schema);
    company_ds.append(reader, None).await.unwrap();
    let company_version = company_ds.version().version;
    let company_metadata =
        table_version_metadata_for_state(uri, &company_entry.table_path, None, company_version)
            .await
            .unwrap();
    mc.commit(&[SubTableUpdate {
        table_key: "node:Company".to_string(),
        table_version: company_version,
        table_branch: None,
        row_count: 1,
        version_metadata: company_metadata,
    }])
    .await
    .unwrap();

    // Writer B writes Person but asserts Company is still at v=1.
    let update_person = append_person_and_make_update(uri, &person_entry, "Bob").await;
    let mut expected = HashMap::new();
    expected.insert("node:Company".to_string(), 1);

    let err = mc
        .commit_with_expected(&[update_person], &expected)
        .await
        .expect_err("drift on an untouched expected table should reject");

    let OmniError::Manifest(m) = err else {
        panic!("expected OmniError::Manifest");
    };
    match m.details {
        Some(ManifestConflictDetails::ExpectedVersionMismatch {
            ref table_key,
            expected,
            actual,
        }) => {
            assert_eq!(table_key, "node:Company");
            assert_eq!(expected, 1);
            assert_eq!(actual, company_version);
        }
        other => panic!("expected ExpectedVersionMismatch, got {:?}", other),
    }
}

#[tokio::test]
async fn test_commit_with_expected_unknown_table_reports_actual_zero() {
    use crate::error::ManifestConflictDetails;

    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();

    let mut expected = HashMap::new();
    expected.insert("node:DoesNotExist".to_string(), 7);
    let err = mc
        .commit_with_expected(&[], &expected)
        .await
        .expect_err("unknown expected table should reject");

    let OmniError::Manifest(m) = err else {
        panic!("expected OmniError::Manifest");
    };
    match m.details {
        Some(ManifestConflictDetails::ExpectedVersionMismatch {
            table_key,
            expected,
            actual,
        }) => {
            assert_eq!(table_key, "node:DoesNotExist");
            assert_eq!(expected, 7);
            assert_eq!(actual, 0);
        }
        other => panic!("expected ExpectedVersionMismatch, got {:?}", other),
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_concurrent_publish_with_overlapping_expected_versions_one_succeeds() {
    use crate::error::ManifestConflictDetails;

    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
    let person_entry = mc.snapshot().entry("node:Person").unwrap().clone();

    // Advance the Person dataset once so we have a real on-disk version 2 that
    // both publishers can target. Both attempt to land the *same*
    // `version:node:Person@v=2` row in `__manifest`, which is the row-level
    // CAS conflict the publisher must detect: load_publish_state at the same
    // baseline → pre-check passes for both → only one merge_insert can land
    // the unique `object_id`.
    let update = append_person_and_make_update(uri, &person_entry, "Alice").await;

    let mut expected = HashMap::new();
    expected.insert("node:Person".to_string(), 1);

    let publisher_a = GraphNamespacePublisher::new(uri, None);
    let publisher_b = GraphNamespacePublisher::new(uri, None);
    let changes_a = vec![ManifestChange::Update(update.clone())];
    let changes_b = vec![ManifestChange::Update(update)];
    let expected_a = expected.clone();
    let expected_b = expected;

    let (res_a, res_b) = tokio::join!(
        async { publisher_a.publish(&changes_a, &expected_a).await },
        async { publisher_b.publish(&changes_b, &expected_b).await }
    );

    let (succeeded, err) = match (res_a, res_b) {
        (Ok(_), Err(e)) => (1, e),
        (Err(e), Ok(_)) => (1, e),
        (Ok(_), Ok(_)) => panic!("both writers committed -- OCC failed"),
        (Err(a), Err(b)) => panic!("both writers failed: {:?} / {:?}", a, b),
    };
    assert_eq!(succeeded, 1, "exactly one writer must succeed");

    let OmniError::Manifest(m) = err else {
        panic!("expected OmniError::Manifest, got {:?}", err);
    };
    // The losing writer surfaces either ExpectedVersionMismatch (its retry's
    // pre-check observed the winner's advance) or a plain Conflict (Lance
    // row-level CAS rejected, retry exhausted before the pre-check fired).
    // Both are acceptable typed conflict signals; what matters is that the
    // failure is not silent.
    use crate::error::ManifestErrorKind;
    assert!(
        matches!(m.kind, ManifestErrorKind::Conflict),
        "expected Conflict-kind manifest error, got {:?}: {}",
        m.kind,
        m.message,
    );
    if let Some(ManifestConflictDetails::ExpectedVersionMismatch {
        ref table_key,
        expected,
        ..
    }) = m.details
    {
        assert_eq!(table_key, "node:Person");
        assert_eq!(expected, 1);
    }

    // Manifest must reflect exactly one new commit on Person at the requested
    // version (no duplicate version rows).
    let mc = ManifestCoordinator::open(uri).await.unwrap();
    let entry = mc.snapshot().entry("node:Person").unwrap().clone();
    assert!(
        entry.table_version > 1,
        "Person should have advanced past v=1"
    );
}

#[tokio::test]
async fn test_init_stamps_internal_schema_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    ManifestCoordinator::init(uri, &catalog).await.unwrap();

    let ds = open_manifest_dataset(uri, None).await.unwrap();
    assert_eq!(
        super::migrations::read_stamp(&ds),
        super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
        "init should stamp the manifest at the current internal schema version",
    );
}

#[tokio::test]
async fn test_publish_migrates_pre_stamp_manifest_to_current_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();

    // Simulate a v1 (pre-stamp) graph by removing the schema-level stamp on disk.
    {
        let mut ds = open_manifest_dataset(uri, None).await.unwrap();
        ds.update_schema_metadata([(
            "omnigraph:internal_schema_version".to_string(),
            None::<String>,
        )])
        .await
        .unwrap();
        let post = open_manifest_dataset(uri, None).await.unwrap();
        assert_eq!(
            super::migrations::read_stamp(&post),
            1,
            "stamp removed ⇒ read_stamp falls back to v1",
        );
    }

    // Publish a no-op (empty changes) but require state to be loaded by passing
    // an expected_table_versions that matches the initial state. This forces
    // the publisher's open-for-write path, which runs the migration.
    let mut expected = HashMap::new();
    expected.insert("node:Person".to_string(), 1);
    GraphNamespacePublisher::new(uri, None)
        .publish(&[], &expected)
        .await
        .unwrap();

    let post = open_manifest_dataset(uri, None).await.unwrap();
    assert_eq!(
        super::migrations::read_stamp(&post),
        super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
        "publish on a v1 graph should leave the manifest stamped at the current version",
    );

    // Manifest should still serve correctly post-migration.
    drop(mc);
    let reopened = ManifestCoordinator::open(uri).await.unwrap();
    assert!(reopened.snapshot().entry("node:Person").is_some());
}

#[tokio::test]
async fn test_v2_to_v3_sweeps_legacy_run_branches_on_write_open() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();

    // Synthesize a pre-MR-770 graph: several stale `__run__` staging branches
    // left on `__manifest` (a real legacy graph accumulates one per run), plus
    // a real user branch that must survive the sweep. Multiple run branches
    // exercise the migration's delete loop on a single reused dataset handle.
    mc.create_branch("__run__01J9LEGACY").await.unwrap();
    mc.create_branch("__run__01J9SECOND").await.unwrap();
    mc.create_branch("__run__01J9THIRD").await.unwrap();
    mc.create_branch("feature").await.unwrap();
    let before = mc.list_branches().await.unwrap();
    assert_eq!(
        before.iter().filter(|b| b.starts_with("__run__")).count(),
        3,
        "precondition: three legacy run branches exist on __manifest; got {before:?}",
    );

    // Rewind the internal-schema stamp to v2 so the next write-open runs the
    // v2 → v3 sweep arm (init stamps at the current version, which is past it).
    {
        let mut ds = open_manifest_dataset(uri, None).await.unwrap();
        ds.update_schema_metadata([(
            "omnigraph:internal_schema_version".to_string(),
            Some("2".to_string()),
        )])
        .await
        .unwrap();
        let post = open_manifest_dataset(uri, None).await.unwrap();
        assert_eq!(super::migrations::read_stamp(&post), 2, "stamp rewound to v2");
    }

    // A no-op publish forces the open-for-write path, which runs the migration.
    let mut expected = HashMap::new();
    expected.insert("node:Person".to_string(), 1);
    GraphNamespacePublisher::new(uri, None)
        .publish(&[], &expected)
        .await
        .unwrap();

    // Stamp advanced to current; the legacy run branch is physically gone from
    // `__manifest` (checked via the raw, unfiltered manifest list — not the
    // guard-filtered `branch_list`), and the real branch + `main` survive.
    let post = open_manifest_dataset(uri, None).await.unwrap();
    assert_eq!(
        super::migrations::read_stamp(&post),
        super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
    );
    let reopened = ManifestCoordinator::open(uri).await.unwrap();
    let after = reopened.list_branches().await.unwrap();
    assert!(
        !after.iter().any(|b| b.starts_with("__run__")),
        "legacy run branch must be swept; got {after:?}",
    );
    assert!(after.iter().any(|b| b == "feature"), "user branch must survive");
    assert!(after.iter().any(|b| b == "main"), "main must survive");

    // Idempotent: a second write-open finds the stamp at current and does not
    // re-run the sweep or error.
    GraphNamespacePublisher::new(uri, None)
        .publish(&[], &expected)
        .await
        .unwrap();
    let final_ds = open_manifest_dataset(uri, None).await.unwrap();
    assert_eq!(
        super::migrations::read_stamp(&final_ds),
        super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
    );
}

#[tokio::test]
async fn test_publish_rejects_manifest_stamped_at_future_version() {
    let dir = tempfile::tempdir().unwrap();
    let uri = dir.path().to_str().unwrap();
    let catalog = build_test_catalog();
    ManifestCoordinator::init(uri, &catalog).await.unwrap();

    // Stamp the manifest at a version higher than this binary knows about.
    let future = super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION + 99;
    {
        let mut ds = open_manifest_dataset(uri, None).await.unwrap();
        ds.update_schema_metadata([(
            "omnigraph:internal_schema_version".to_string(),
            Some(future.to_string()),
        )])
        .await
        .unwrap();
    }

    let mut expected = HashMap::new();
    expected.insert("node:Person".to_string(), 1);
    let err = GraphNamespacePublisher::new(uri, None)
        .publish(&[], &expected)
        .await
        .expect_err("future-stamped manifest should reject open-for-write");
    let msg = err.to_string();
    assert!(
        msg.contains("upgrade omnigraph") && msg.contains(&future.to_string()),
        "expected forward-version refusal, got: {}",
        msg,
    );
}

#[test]
fn manifest_column_helpers_return_error_for_bad_schema() {
    let batch = RecordBatch::try_new(
        Arc::new(Schema::new(vec![Field::new(
            "table_key",
            DataType::UInt64,
            false,
        )])),
        vec![Arc::new(UInt64Array::from(vec![1_u64]))],
    )
    .unwrap();

    let err = string_column(&batch, "table_key").unwrap_err();
    assert!(err.to_string().contains("table_key"));
}
