genegraph-storage 0.72.0

use std::collections::{BTreeMap, HashMap};
use std::path::Path;

use crate::catalog::{Catalog, CollectionKind, GraphStaleness, LocalRegistry, TableDescriptor};
use crate::lance_storage_graph::LanceStorageGraph;
use crate::metadata::{FileInfo, GeneMetadata};
use crate::tests::tmp_dir;
use crate::traits::metadata::Metadata;

fn registry_with(base: &Path, name_id: &str) -> LocalRegistry {
    let metadata = GeneMetadata::new(name_id)
        .with_base(base.to_path_buf())
        .with_dimensions(100, 50);
    let raw_info = metadata
        .new_fileinfo("rawinput", "dense", (100, 50), None, None)
        .expect("dense filetype");
    let adj_info = metadata
        .new_fileinfo("adjacency", "sparse", (100, 100), Some(500), None)
        .expect("sparse filetype");
    let metadata = metadata
        .add_file("rawinput", raw_info)
        .add_file("adjacency", adj_info);
    LocalRegistry::new(metadata, base.to_path_buf())
}

#[tokio::test(flavor = "multi_thread")]
async fn catalog_lists_and_describes_tables() {
    let base = tmp_dir("catalog_m_c1").await;
    let name_id = "catalog_test";
    let registry = registry_with(&base, name_id);

    let tables = registry.list_tables().expect("list");
    let names: Vec<_> = tables.iter().map(|t| t.name.clone()).collect();
    assert_eq!(names, vec!["adjacency", "rawinput"], "sorted by name");
    assert!(tables.iter().all(|t| t.format == "lance"));

    let adj = registry.describe_table("adjacency").expect("describe");
    assert_eq!(
        adj.base_location,
        base.join(format!("{name_id}_adjacency.lance"))
    );
    assert_eq!(
        adj.properties.get("filetype").map(String::as_str),
        Some("sparse")
    );
    assert_eq!(adj.properties.get("nnz").map(String::as_str), Some("500"));

    assert!(registry.table_exists("rawinput").unwrap());
    assert!(!registry.table_exists("missing").unwrap());
    assert!(registry.describe_table("missing").is_err());
}

#[tokio::test(flavor = "multi_thread")]
async fn catalog_register_and_deregister_roundtrip() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "catalog_test");

    registry
        .register_table(TableDescriptor {
            name: "norms".to_string(),
            format: "lance".to_string(),
            base_location: base.join("catalog_test_norms.lance"),
            kind: CollectionKind::VectorSpace,
            properties: HashMap::from([
                ("filetype".to_string(), "vector".to_string()),
                ("rows".to_string(), "100".to_string()),
                ("cols".to_string(), "1".to_string()),
            ])
            .into_iter()
            .collect(),
        })
        .expect("register");

    assert!(registry.table_exists("norms").unwrap());
    let norms = registry.describe_table("norms").unwrap();
    assert_eq!(norms.base_location, base.join("catalog_test_norms.lance"));
    assert_eq!(
        norms.properties.get("rows").map(String::as_str),
        Some("100")
    );

    registry.deregister_table("norms").expect("deregister");
    assert!(!registry.table_exists("norms").unwrap());
    assert!(registry.deregister_table("norms").is_err());

    // the mutated metadata must be recoverable for persistence
    let metadata = registry.into_metadata();
    assert!(!metadata.files.contains_key("norms"));
    assert!(metadata.files.contains_key("rawinput"));
}

#[tokio::test(flavor = "multi_thread")]
async fn catalog_rejects_non_lance_format_and_bad_properties() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "catalog_test");

    let err = registry
        .register_table(TableDescriptor {
            name: "delta_thing".to_string(),
            format: "delta".to_string(),
            base_location: base.join("delta_thing"),
            kind: CollectionKind::Table,
            properties: Default::default(),
        })
        .unwrap_err();
    assert!(matches!(err, crate::StorageError::UnsupportedFormat(_)));

    let err = registry
        .register_table(TableDescriptor {
            name: "bad_props".to_string(),
            format: "lance".to_string(),
            base_location: base.join("bad_props.lance"),
            kind: CollectionKind::Table,
            properties: HashMap::from([("rows".to_string(), "many".to_string())])
                .into_iter()
                .collect(),
        })
        .unwrap_err();
    assert!(matches!(err, crate::StorageError::Invalid(_)));
}

/// RFC #81-P1: legacy artifact keys become pre-seeded collections — kinds are
/// derived from the filetype (dense/vector -> vector-space, sparse -> graph)
/// and descriptors carry a `kind` property alongside the typed field.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_derives_kinds_for_legacy_keys() {
    let base = tmp_dir("catalog_m_c1").await;
    let registry = registry_with(&base, "catalog_test");

    let raw = registry.describe_table("rawinput").unwrap();
    assert_eq!(raw.kind, CollectionKind::VectorSpace);
    assert_eq!(
        raw.properties.get("kind").map(String::as_str),
        Some("vector-space")
    );

    let adj = registry.describe_table("adjacency").unwrap();
    assert_eq!(adj.kind, CollectionKind::Graph);
    assert_eq!(
        adj.properties.get("kind").map(String::as_str),
        Some("graph")
    );
}

/// RFC #81-P1: user properties and an explicit kind round-trip through
/// `register_table` -> `describe_table`, and are persisted inside the
/// metadata `files` map (registry-level kinds).
#[tokio::test(flavor = "multi_thread")]
async fn catalog_kinds_and_user_properties_roundtrip() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "catalog_test");

    registry
        .register_table(TableDescriptor {
            name: "edges".to_string(),
            format: "lance".to_string(),
            base_location: base.join("catalog_test_edges.lance"),
            kind: CollectionKind::Graph,
            properties: BTreeMap::from([
                ("filetype".to_string(), "graph".to_string()),
                ("rows".to_string(), "7".to_string()),
                ("cols".to_string(), "7".to_string()),
                ("nnz".to_string(), "12".to_string()),
                ("node_id_width".to_string(), "u32".to_string()),
            ]),
        })
        .expect("register graph collection");

    let desc = registry.describe_table("edges").unwrap();
    assert_eq!(desc.kind, CollectionKind::Graph);
    assert_eq!(
        desc.properties.get("node_id_width").map(String::as_str),
        Some("u32"),
        "user properties pass through"
    );

    // the user property must be stored on the FileInfo itself
    let metadata = registry.into_metadata();
    let info = metadata.files.get("edges").unwrap();
    assert_eq!(info.kind, Some(CollectionKind::Graph));
    assert_eq!(
        info.properties.get("node_id_width").map(String::as_str),
        Some("u32")
    );
}

/// RFC #81-P1: pre-0.28 metadata JSON (no `kind` / `properties` fields on
/// FileInfo) still deserializes and the registry derives collection kinds.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_legacy_metadata_json_still_loads() {
    let legacy_json = r#"{
        "name_id": "legacy_instance",
        "nrows": 10,
        "ncols": 5,
        "base": "/tmp/legacy_instance",
        "files": {
            "rawinput": {
                "filename": "legacy_instance_rawinput.lance",
                "filetype": "dense",
                "storage_format": "lance fixed-row",
                "rows": 10,
                "cols": 5,
                "nnz": null,
                "size_bytes": null
            }
        },
        "created_at": "2026-01-01T00:00:00Z"
    }"#;
    let metadata: GeneMetadata = serde_json::from_str(legacy_json).expect("legacy JSON loads");
    assert!(metadata.files["rawinput"].kind.is_none());
    assert!(metadata.files["rawinput"].properties.is_empty());

    let registry = LocalRegistry::new(metadata, std::path::PathBuf::from("/tmp/legacy_instance"));
    let desc = registry.describe_table("rawinput").unwrap();
    assert_eq!(desc.kind, CollectionKind::VectorSpace);
}

/// RFC #81-P4: `describe_vector_space` returns the vector-space descriptor
/// plus the graph collection referenced by the `graph` user property.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_describe_vector_space_links_graph() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "catalog_test");

    registry
        .register_table(TableDescriptor {
            name: "embeddings".to_string(),
            format: "lance".to_string(),
            base_location: base.join("catalog_test_embeddings.lance"),
            kind: CollectionKind::VectorSpace,
            properties: BTreeMap::from([
                ("filetype".to_string(), "vectors".to_string()),
                ("rows".to_string(), "100".to_string()),
                ("cols".to_string(), "50".to_string()),
                ("graph".to_string(), "adjacency".to_string()),
            ]),
        })
        .expect("register vector space");

    let vs = registry.describe_vector_space("embeddings").unwrap();
    assert_eq!(vs.vectors.name, "embeddings");
    let graph = vs.graph.expect("linked graph");
    assert_eq!(graph.name, "adjacency");
    assert_eq!(graph.kind, CollectionKind::Graph);

    // a vector space without a graph link has no linked descriptor
    registry
        .register_table(TableDescriptor {
            name: "bare".to_string(),
            format: "lance".to_string(),
            base_location: base.join("catalog_test_bare.lance"),
            kind: CollectionKind::VectorSpace,
            properties: BTreeMap::from([
                ("filetype".to_string(), "vectors".to_string()),
                ("rows".to_string(), "4".to_string()),
                ("cols".to_string(), "4".to_string()),
            ]),
        })
        .expect("register bare vector space");
    assert!(
        registry
            .describe_vector_space("bare")
            .unwrap()
            .graph
            .is_none()
    );

    // kind mismatches and dangling links are errors, not guesses
    let err = registry.describe_vector_space("adjacency").unwrap_err();
    assert!(matches!(err, crate::StorageError::Invalid(_)));
    assert!(registry.describe_vector_space("missing").is_err());

    registry
        .register_table(TableDescriptor {
            name: "dangling".to_string(),
            format: "lance".to_string(),
            base_location: base.join("catalog_test_dangling.lance"),
            kind: CollectionKind::VectorSpace,
            properties: BTreeMap::from([
                ("filetype".to_string(), "vectors".to_string()),
                ("graph".to_string(), "no_such_graph".to_string()),
            ]),
        })
        .expect("register dangling vector space");
    let err = registry.describe_vector_space("dangling").unwrap_err();
    assert!(matches!(err, crate::StorageError::Invalid(_)));
}

/// The registry resolves against a real LanceStorageGraph dataset: locations
/// point at the actual `*.lance` directories.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_locations_match_storage_layout() {
    let base = tmp_dir("catalog_m_c1_storage").await;
    let name_id = "catalog_storage";
    let storage = LanceStorageGraph::new(base.to_string_lossy().to_string(), name_id.to_string())
        .expect("valid instance name");
    let metadata = GeneMetadata::seed_metadata(name_id, 10, 5, &storage)
        .await
        .expect("seed");

    let registry = LocalRegistry::new(metadata, base.to_path_buf());
    assert!(
        registry.list_tables().unwrap().is_empty(),
        "fresh dataset has no tables"
    );
    assert_eq!(
        registry.describe_table("rawinput").unwrap_err().to_string(),
        "Invalid data: table 'rawinput' is not registered"
    );
}

// ---------------------------------------------------------------------------
// #140: graph source lineage + vector-space staleness
// ---------------------------------------------------------------------------

/// #140: the catalog resolves a graph source against the registry. The
/// source must be a registered vector-space collection; the resolved
/// lineage is what a graph write stamps as `source`/`source_rows`.
#[tokio::test(flavor = "multi_thread")]
async fn graph_source_resolves_only_registered_vector_spaces() {
    let base = tmp_dir("catalog_m_c1").await;
    let registry = registry_with(&base, "graph_source");

    // rawinput: dense -> vector space with 100 rows
    let src = registry.graph_source("rawinput").expect("resolves");
    assert_eq!(src.name, "rawinput");
    assert_eq!(src.rows, 100);

    // adjacency: sparse -> graph kind is not a valid source
    let err = registry.graph_source("adjacency").unwrap_err();
    assert!(
        matches!(err, crate::StorageError::Invalid(_)),
        "expected Invalid for graph-kind source, got {err:?}"
    );

    let err = registry.graph_source("missing").unwrap_err();
    assert!(
        matches!(err, crate::StorageError::Invalid(_)),
        "expected Invalid for unregistered source, got {err:?}"
    );
}

/// #140: `describe_vector_space` computes graph staleness from current
/// registry facts. Lineage (`source`/`source_rows`) is authoritative when
/// present; without it a `num_nodes`/rows mismatch is `Stale`, while an
/// agreement is `Unknown` — counts agree, content may differ, no guesses.
#[tokio::test(flavor = "multi_thread")]
async fn describe_vector_space_reports_graph_staleness_states() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "staleness_states");

    let vector_space = |name: &str, extra: &[(&str, &str)]| TableDescriptor {
        name: name.to_string(),
        format: "lance".to_string(),
        base_location: base.join(format!("staleness_states_{name}.lance")),
        kind: CollectionKind::VectorSpace,
        properties: {
            let mut props = BTreeMap::from([
                ("filetype".to_string(), "vectors".to_string()),
                ("rows".to_string(), "100".to_string()),
                ("cols".to_string(), "4".to_string()),
            ]);
            props.extend(extra.iter().map(|(k, v)| (k.to_string(), v.to_string())));
            props
        },
    };
    let graph = |name: &str, extra: &[(&str, &str)]| TableDescriptor {
        name: name.to_string(),
        format: "lance".to_string(),
        base_location: base.join(format!("staleness_states_{name}.lance")),
        kind: CollectionKind::Graph,
        properties: {
            let mut props = BTreeMap::from([
                ("filetype".to_string(), "graph".to_string()),
                ("rows".to_string(), "200".to_string()),
                ("cols".to_string(), "3".to_string()),
            ]);
            props.extend(extra.iter().map(|(k, v)| (k.to_string(), v.to_string())));
            props
        },
    };

    // unlinked: no `graph` property, nothing to assess
    registry.register_table(vector_space("bare", &[])).unwrap();
    assert_eq!(
        registry
            .describe_vector_space("bare")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Unlinked
    );

    // fresh: lineage present and matching this vector space
    registry
        .register_table(graph(
            "g_fresh",
            &[
                ("source", "vs_fresh"),
                ("source_rows", "100"),
                ("num_nodes", "100"),
            ],
        ))
        .unwrap();
    registry
        .register_table(vector_space("vs_fresh", &[("graph", "g_fresh")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_fresh")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Fresh
    );

    // stale: the lineage names another collection
    registry
        .register_table(graph(
            "g_foreign",
            &[("source", "someone_else"), ("source_rows", "100")],
        ))
        .unwrap();
    registry
        .register_table(vector_space("vs_foreign", &[("graph", "g_foreign")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_foreign")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Stale
    );

    // stale: source_rows differs from the current row count
    registry
        .register_table(graph(
            "g_oldrows",
            &[("source", "vs_oldrows"), ("source_rows", "99")],
        ))
        .unwrap();
    registry
        .register_table(vector_space("vs_oldrows", &[("graph", "g_oldrows")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_oldrows")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Stale
    );

    // stale without lineage: num_nodes differs from the row count
    registry
        .register_table(graph("g_big", &[("num_nodes", "150")]))
        .unwrap();
    registry
        .register_table(vector_space("vs_big", &[("graph", "g_big")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_big")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Stale
    );

    // unknown without lineage: num_nodes matches rows, content unverified
    registry
        .register_table(graph("g_match", &[("num_nodes", "100")]))
        .unwrap();
    registry
        .register_table(vector_space("vs_match", &[("graph", "g_match")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_match")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Unknown
    );

    // lineage wins over the num_nodes fallback
    registry
        .register_table(graph(
            "g_wins",
            &[
                ("num_nodes", "7"),
                ("source", "vs_wins"),
                ("source_rows", "100"),
            ],
        ))
        .unwrap();
    registry
        .register_table(vector_space("vs_wins", &[("graph", "g_wins")]))
        .unwrap();
    assert_eq!(
        registry
            .describe_vector_space("vs_wins")
            .unwrap()
            .graph_staleness,
        GraphStaleness::Fresh
    );
}

/// #140: the staleness signal reads registry facts only. It computes
/// identically over a zarr-shaped registry (`storage_format = "zzarr"`,
/// rel-path keys, filetypes `dense`/`graph`) — one rule, every format.
#[tokio::test(flavor = "multi_thread")]
async fn staleness_signal_is_storage_format_agnostic() {
    let base = tmp_dir("catalog_m_c1").await;

    let metadata = GeneMetadata::new("zzarr_stale")
        .with_base(base.clone())
        .with_dimensions(100, 50);
    let mut vectors = FileInfo::new("main/vectors".to_string(), "dense", (100, 4), None, None)
        .expect("dense filetype");
    vectors.storage_format = "zzarr".to_string();
    vectors
        .properties
        .insert("graph".to_string(), "main/g".to_string());
    let mut graph = FileInfo::new("main/g".to_string(), "graph", (100, 3), Some(200), None)
        .expect("graph filetype");
    graph.storage_format = "zzarr".to_string();
    graph
        .properties
        .insert("source".to_string(), "main/vectors".to_string());
    graph
        .properties
        .insert("source_rows".to_string(), "100".to_string());
    let metadata = metadata
        .add_file("main/vectors", vectors)
        .add_file("main/g", graph);

    let registry = LocalRegistry::new(metadata, base);
    let desc = registry
        .describe_vector_space("main/vectors")
        .expect("describe zarr-shaped registry");
    assert_eq!(
        desc.vectors
            .properties
            .get("storage_format")
            .map(String::as_str),
        Some("zzarr")
    );
    assert_eq!(desc.graph_staleness, GraphStaleness::Fresh);
}

// ---------------------------------------------------------------------------
// #142: Arrow-IPC artifacts in the catalog
// ---------------------------------------------------------------------------

/// An IPC artifact registered with `storage_format = "arrow-ipc"` describes
/// as an `arrow-ipc` table of the vector-space kind — the semantic kind
/// stays `CollectionKind`, the table format names the storage format.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_describes_arrow_ipc_entries() {
    let base = tmp_dir("catalog_m_c1").await;
    let metadata = GeneMetadata::new("ipc_catalog")
        .with_base(base.clone())
        .with_dimensions(4, 8);
    let mut vectors = FileInfo::new(
        "ipc_catalog_vectors.arrow".to_string(),
        "dense",
        (4, 8),
        None,
        None,
    )
    .expect("dense filetype");
    vectors.storage_format = "arrow-ipc".to_string();
    let registry = LocalRegistry::new(metadata.add_file("vectors", vectors), base);

    let desc = registry
        .describe_table("vectors")
        .expect("describe IPC entry");
    assert_eq!(desc.format, "arrow-ipc");
    assert_eq!(desc.kind, CollectionKind::VectorSpace);
    assert_eq!(
        desc.properties.get("storage_format").map(String::as_str),
        Some("arrow-ipc")
    );
}

/// Streaming IPC artifacts (`arrow-ipc-stream`) register through
/// `register_table` and keep their format through describe; unrelated
/// formats stay rejected.
#[tokio::test(flavor = "multi_thread")]
async fn catalog_registers_arrow_ipc_descriptors() {
    let base = tmp_dir("catalog_m_c1").await;
    let mut registry = registry_with(&base, "ipc_stream_catalog");

    let base_location = base.join("ipc_stream_catalog_stream.arrows");
    registry
        .register_table(TableDescriptor {
            name: "stream".to_string(),
            format: "arrow-ipc-stream".to_string(),
            base_location: base_location.clone(),
            kind: CollectionKind::VectorSpace,
            properties: BTreeMap::from([
                ("filetype".to_string(), "dense".to_string()),
                ("rows".to_string(), "12".to_string()),
                ("cols".to_string(), "4".to_string()),
            ])
            .into_iter()
            .collect(),
        })
        .expect("register arrow-ipc-stream descriptor");

    let desc = registry.describe_table("stream").expect("describe stream");
    assert_eq!(desc.format, "arrow-ipc-stream");
    assert_eq!(desc.base_location, base_location);
    // the declared format survives into the registry entry
    assert_eq!(
        registry
            .metadata()
            .files
            .get("stream")
            .expect("registry entry")
            .storage_format,
        "arrow-ipc-stream"
    );

    // unrelated formats keep their typed rejection
    let err = registry
        .register_table(TableDescriptor {
            name: "delta_thing".to_string(),
            format: "delta".to_string(),
            base_location: base.join("delta_thing"),
            kind: CollectionKind::Table,
            properties: Default::default(),
        })
        .unwrap_err();
    assert!(matches!(err, crate::StorageError::UnsupportedFormat(_)));
}