linkrs-storage 0.1.0

use std::sync::Arc;

use crate::engine::graph_storage::context::GraphStorageContext;
use crate::index::manifest::{GenerationBuildState, GenerationState};
use crate::index::types::IndexRecord;
use crate::{GraphStorage, StoragePersistenceOps, StorageReader, StorageSchemaOps, StorageWriter};
use linkrs_core::types::{
    CommitLsn, IdempotencyKey, Index, IndexConfig, IndexField, IndexGeneration, IndexType,
    OrderingKey, SnapshotTimestamp, TargetId, TransactionId, VertexId,
};
use linkrs_core::wal::{EntityRef, IndexMutation, IndexOperation, OutboxIntent};
use linkrs_core::Value;

use super::checkpoint::load_generation_build_state;
use super::checkpoint::{
    build_vertex_index_data, remove_generation_build_state, resolve_crash_recovery,
    save_generation_build_state, write_generation_checkpoint,
};
use super::generation::GenerationFaultPoint;
use super::wal_replay::replay_wal_partition;

fn setup_context() -> GraphStorageContext {
    GraphStorageContext::new()
}

fn generation_state(generation: u64, start_lsn: u64) -> GenerationBuildState {
    GenerationBuildState::new(
        IndexGeneration::new(generation),
        SnapshotTimestamp::new(start_lsn),
        CommitLsn::new(start_lsn),
    )
}

fn setup_persistent_context() -> (tempfile::TempDir, GraphStorageContext) {
    let temp_dir = tempfile::TempDir::new().expect("Failed to create temp dir");
    let config = crate::engine::PersistenceConfig::for_work_dir(temp_dir.path());
    let ctx = GraphStorageContext::new_with_persistence(temp_dir.path().to_path_buf(), config)
        .expect("Failed to create persistent context");
    (temp_dir, ctx)
}

fn test_intent(transaction_id: TransactionId, index_id: u64, vertex_id: i64) -> OutboxIntent {
    OutboxIntent {
        wire_version: linkrs_core::wal::WAL_SYNC_WIRE_VERSION,
        transaction_id,
        intent_sequence: 0,
        mutation: IndexMutation {
            wire_version: linkrs_core::wal::WAL_SYNC_WIRE_VERSION,
            target: TargetId::new("native-index").expect("target should be valid"),
            index_id,
            index_generation: IndexGeneration::new(1),
            entity_ref: EntityRef::Vertex(
                VertexId::try_from_int64(vertex_id).expect("test vertex id"),
            ),
            operation: IndexOperation::Upsert,
            document_or_vector: Vec::new(),
            idempotency_key: IdempotencyKey::new(format!("intent-{transaction_id}-{vertex_id}"))
                .expect("idempotency key should be valid"),
            ordering_key: OrderingKey::new(format!("vertex-{vertex_id}"))
                .expect("ordering key should be valid"),
        },
    }
}

#[test]
fn wal_catch_up_uses_only_intent_entities() {
    let entity = EntityRef::Vertex(VertexId::try_from_int64(1).expect("test vertex id"));
    let other_entity = EntityRef::Vertex(VertexId::try_from_int64(2).expect("test vertex id"));
    let intents = vec![test_intent(TransactionId::new(1), 7, 1)];

    let mut active_forward = std::collections::BTreeMap::new();
    active_forward.insert(
        b"logged-forward".to_vec(),
        IndexRecord::new(20).with_entity_ref(entity.clone()),
    );
    active_forward.insert(
        b"memory-only-forward".to_vec(),
        IndexRecord::new(20).with_entity_ref(other_entity.clone()),
    );
    let mut active_reverse = std::collections::BTreeMap::new();
    active_reverse.insert(
        b"logged-reverse".to_vec(),
        IndexRecord::new(20).with_entity_ref(entity),
    );
    active_reverse.insert(
        b"memory-only-reverse".to_vec(),
        IndexRecord::new(20).with_entity_ref(other_entity),
    );

    let mut rebuilt_forward = std::collections::BTreeMap::new();
    rebuilt_forward.insert(b"snapshot-forward".to_vec(), IndexRecord::new(10));
    let mut rebuilt_reverse = std::collections::BTreeMap::new();
    rebuilt_reverse.insert(b"snapshot-reverse".to_vec(), IndexRecord::new(10));

    let (forward, reverse) = replay_wal_partition(
        (active_forward, active_reverse),
        (rebuilt_forward, rebuilt_reverse),
        10,
        &intents,
        |_| true,
        |_| true,
    );

    assert!(forward.contains_key(b"snapshot-forward".as_slice()));
    assert!(forward.contains_key(b"logged-forward".as_slice()));
    assert!(!forward.contains_key(b"memory-only-forward".as_slice()));
    assert!(reverse.contains_key(b"snapshot-reverse".as_slice()));
    assert!(reverse.contains_key(b"logged-reverse".as_slice()));
    assert!(!reverse.contains_key(b"memory-only-reverse".as_slice()));
}

#[test]
fn wal_catch_up_reads_committed_intents_from_disk() {
    let (_temp_dir, ctx) = setup_persistent_context();
    let transaction_id = TransactionId::new(9);
    let intent = test_intent(transaction_id, 7, 42);

    ctx.commit_staged_writes(transaction_id, &[intent])
        .expect("transaction should be durable");
    let barrier_lsn = super::current_wal_lsn(&ctx);
    let index = Index::new(IndexConfig {
        id: 7,
        name: "test_index".to_string(),
        space_id: 1,
        schema_name: "Person".to_string(),
        fields: vec![IndexField::new(
            "name".to_string(),
            Value::string(""),
            false,
        )],
        properties: Vec::new(),
        index_type: IndexType::TagIndex,
        is_unique: false,
        covering: false,
        partial_condition: None,
    });
    let filtered = super::wal_intents_for_index(&ctx, 1, &index, CommitLsn::ZERO, barrier_lsn)
        .expect("committed intents should be readable");

    assert_eq!(filtered.len(), 1);
    assert_eq!(filtered[0].transaction_id, transaction_id);
    assert_eq!(filtered[0].mutation.index_id, 7);
}

#[test]
fn concurrent_rebuild_and_writes_preserve_new_index_entries() {
    use std::collections::HashMap;
    use std::thread;

    let temp_dir = tempfile::TempDir::new().expect("temporary directory should be created");
    let mut storage = GraphStorage::new_with_path(temp_dir.path().to_path_buf())
        .expect("persistent storage should be created");
    let mut space = linkrs_core::types::SpaceInfo::new("test_space".to_string())
        .with_vid_type(linkrs_core::DataType::BigInt);
    storage
        .create_space(&mut space)
        .expect("space should be created");
    let tag = linkrs_core::types::TagInfo::new("Person".to_string()).with_properties(vec![
        linkrs_core::types::PropertyDef::new("id".into(), linkrs_core::DataType::BigInt),
        linkrs_core::types::PropertyDef::new("name".into(), linkrs_core::DataType::String),
    ]);
    storage
        .create_tag("test_space", &tag)
        .expect("tag should be created");
    let index = Index::new(IndexConfig {
        id: 1,
        name: "person_name_idx".to_string(),
        space_id: 1,
        schema_name: "Person".to_string(),
        fields: vec![IndexField::new(
            "name".to_string(),
            Value::string(""),
            false,
        )],
        properties: Vec::new(),
        index_type: IndexType::TagIndex,
        is_unique: false,
        covering: false,
        partial_condition: None,
    });
    storage
        .create_tag_index("test_space", &index)
        .expect("index should be created");

    for vertex_id in 0..500 {
        let mut properties: HashMap<Arc<str>, Value> = HashMap::new();
        properties.insert(
            Arc::from("name"),
            Value::string(format!("initial-{vertex_id}")),
        );
        storage
            .insert_vertex(
                "test_space",
                linkrs_core::Vertex::new(
                    VertexId::try_from_int64(vertex_id).expect("test vertex id"),
                    linkrs_core::vertex_edge_path::Tag::new("Person".into(), properties),
                ),
            )
            .expect("initial vertex should be inserted");
    }

    let rebuild_storage = storage.clone();
    let writer_storage = storage.clone();
    let rebuild = thread::spawn(move || {
        let mut rebuild_storage = rebuild_storage;
        rebuild_storage
            .rebuild_tag_index("test_space", "person_name_idx")
            .expect("rebuild should succeed");
    });
    let writer = thread::spawn(move || {
        let mut writer_storage = writer_storage;
        for vertex_id in 10_000..10_050 {
            let mut properties: HashMap<Arc<str>, Value> = HashMap::new();
            properties.insert(
                Arc::from("name"),
                Value::string(format!("concurrent-{vertex_id}")),
            );
            writer_storage
                .insert_vertex(
                    "test_space",
                    linkrs_core::Vertex::new(
                        VertexId::try_from_int64(vertex_id).expect("test vertex id"),
                        linkrs_core::vertex_edge_path::Tag::new("Person".into(), properties),
                    ),
                )
                .expect("concurrent vertex should be inserted");
        }
    });

    rebuild.join().expect("rebuild thread should finish");
    writer.join().expect("writer thread should finish");

    for vertex_id in 10_000..10_050 {
        let indexed = storage
            .lookup_index(
                "test_space",
                "person_name_idx",
                &Value::string(format!("concurrent-{vertex_id}")),
            )
            .expect("index lookup should succeed");
        assert_eq!(
            indexed,
            vec![Value::from(
                VertexId::try_from_int64(vertex_id).expect("test vertex id")
            )]
        );
    }
}

#[test]
fn rebuild_restarts_after_incremental_replay_failure() {
    let temp_dir = tempfile::TempDir::new().expect("temporary directory should be created");
    let index = Index::new(IndexConfig {
        id: 1,
        name: "person_name_idx".to_string(),
        space_id: 1,
        schema_name: "Person".to_string(),
        fields: vec![IndexField::new(
            "name".to_string(),
            Value::string(""),
            false,
        )],
        properties: Vec::new(),
        index_type: IndexType::TagIndex,
        is_unique: false,
        covering: false,
        partial_condition: None,
    });

    {
        let mut storage = GraphStorage::new_with_path(temp_dir.path().to_path_buf())
            .expect("persistent storage should be created");
        let mut space = linkrs_core::types::SpaceInfo::new("test_space".to_string())
            .with_vid_type(linkrs_core::DataType::BigInt);
        storage
            .create_space(&mut space)
            .expect("space should be created");
        let tag = linkrs_core::types::TagInfo::new("Person".to_string()).with_properties(vec![
            linkrs_core::types::PropertyDef::new("id".into(), linkrs_core::DataType::BigInt),
            linkrs_core::types::PropertyDef::new("name".to_string(), linkrs_core::DataType::String),
        ]);
        storage
            .create_tag("test_space", &tag)
            .expect("tag should be created");
        storage
            .create_tag_index("test_space", &index)
            .expect("index should be created");
        let vertex = linkrs_core::Vertex::new(
            VertexId::try_from_int64(1).expect("test vertex id"),
            linkrs_core::vertex_edge_path::Tag::new(
                "Person".to_string(),
                vec![("name".into(), Value::string("Alice"))]
                    .into_iter()
                    .collect(),
            ),
        );
        storage
            .insert_vertex("test_space", vertex)
            .expect("vertex should be inserted");
        storage
            .rebuild_tag_index("test_space", "person_name_idx")
            .expect("initial generation should build");
        storage.flush().expect("initial index state should flush");
        storage
            .create_checkpoint()
            .expect("initial checkpoint should succeed");

        super::inject_generation_fault(GenerationFaultPoint::IncrementalReplay);
        let result = storage.rebuild_tag_index("test_space", "person_name_idx");
        super::clear_generation_faults();
        assert!(
            result.is_err(),
            "injected catch-up failure should be returned"
        );
        storage
            .flush()
            .expect("storage should flush before restart");
    }

    let mut storage = GraphStorage::open(temp_dir.path().to_path_buf())
        .expect("storage should reopen after the failed build");
    assert!(storage
        .rebuild_tag_index("test_space", "person_name_idx")
        .expect("rebuild should restart after crash recovery"));
    let indexed = storage
        .lookup_index("test_space", "person_name_idx", &Value::string("Alice"))
        .expect("rebuilt index should be readable");
    assert_eq!(
        indexed,
        vec![Value::from(
            VertexId::try_from_int64(1).expect("test vertex id")
        )]
    );
}

#[test]
fn test_create_and_list_tag_index() {
    let ctx = setup_context();

    let index = Index::new(IndexConfig {
        id: 1,
        name: "person_name_idx".to_string(),
        space_id: 0,
        schema_name: "Person".to_string(),
        fields: vec![IndexField::new(
            "name".to_string(),
            Value::string(""),
            false,
        )],
        properties: vec![],
        index_type: IndexType::TagIndex,
        is_unique: false,
        covering: false,
        partial_condition: None,
    });

    super::create_tag_index(&ctx, "test_space", &index)
        .expect_err("should fail because space does not exist");

    // Actually we need a space + schema adapter for full testing.
    // The index_manager functions require a schema_manager with registered space.
    // This is tested through integration tests (tests.rs).
}

#[test]
fn test_get_tag_index_on_empty() {
    let ctx = setup_context();
    let result = super::get_tag_index(&ctx, "nonexistent", "some_index");
    assert!(result.is_err());
}

#[test]
fn test_drop_tag_index_on_empty() {
    let ctx = setup_context();
    let result = super::drop_tag_index(&ctx, "nonexistent", "some_index");
    assert!(result.is_err());
}

#[test]
fn test_lookup_index_on_nonexistent_space() {
    let ctx = setup_context();
    let result = super::lookup_index(&ctx, "no_space", "some_index", &Value::string("test"));
    assert!(result.is_err());
}

// ── Generation build state persistence tests ───────────────────────────

#[test]
fn test_generation_build_state_save_load_remove_roundtrip() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "test_idx";

    let state = generation_state(42, 100);

    // Save
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("build state should save");
    assert_eq!(state.state, GenerationState::Building);
    assert_eq!(state.generation, IndexGeneration::new(42));
    assert_eq!(state.start_lsn, CommitLsn::new(100));
    assert!(state.barrier_lsn.is_none());

    // Load
    let loaded = load_generation_build_state(&ctx, space_id, index_name)
        .expect("build state should load")
        .expect("build state should exist");
    assert_eq!(loaded.generation, state.generation);
    assert_eq!(loaded.start_lsn, state.start_lsn);
    assert_eq!(loaded.state, state.state);
    assert!(loaded.barrier_lsn.is_none());

    // Remove
    remove_generation_build_state(&ctx, space_id, index_name).expect("build state should remove");
    let after_remove =
        load_generation_build_state(&ctx, space_id, index_name).expect("load should succeed");
    assert!(after_remove.is_none());
}

#[test]
fn test_generation_build_state_transitions_are_persistent() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "transition_idx";

    let mut state = generation_state(1, 10);
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("Building state should save");

    // Transition to CatchingUp
    state
        .transition_to_catching_up()
        .expect("Building should transition to CatchingUp");
    assert_eq!(state.state, GenerationState::CatchingUp);
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("CatchingUp state should save");

    let loaded = load_generation_build_state(&ctx, space_id, index_name)
        .expect("load should succeed")
        .expect("build state should exist");
    assert_eq!(loaded.state, GenerationState::CatchingUp);

    // Transition to Publishing
    state
        .transition_to_publishing(CommitLsn::new(50))
        .expect("CatchingUp should transition to Publishing");
    assert_eq!(state.state, GenerationState::Publishing);
    assert_eq!(state.barrier_lsn, Some(CommitLsn::new(50)));
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("Publishing state should save");

    let loaded = load_generation_build_state(&ctx, space_id, index_name)
        .expect("load should succeed")
        .expect("build state should exist");
    assert_eq!(loaded.state, GenerationState::Publishing);
    assert_eq!(loaded.barrier_lsn, Some(CommitLsn::new(50)));

    // Transition to Active
    state
        .transition_to_active()
        .expect("Publishing should transition to Active");
    assert!(state.is_active());
}

#[test]
fn test_crash_recovery_discards_incomplete_building_state() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "building_crash_idx";

    // Simulate crash during Building: save Building state then recover
    let state = generation_state(1, 10);
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("Building state should save");

    // Recover (should discard incomplete build)
    resolve_crash_recovery(&ctx, space_id, index_name).expect("recovery should succeed");

    let after =
        load_generation_build_state(&ctx, space_id, index_name).expect("load should succeed");
    assert!(
        after.is_none(),
        "Building state should be discarded on crash recovery"
    );
}

#[test]
fn test_crash_recovery_discards_incomplete_catching_up_state() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "catchup_crash_idx";

    // Simulate crash during CatchingUp
    let mut state = generation_state(1, 10);
    state
        .transition_to_catching_up()
        .expect("Building should transition to CatchingUp");
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("CatchingUp state should save");

    // Recover (should discard incomplete catch-up)
    resolve_crash_recovery(&ctx, space_id, index_name).expect("recovery should succeed");

    let after =
        load_generation_build_state(&ctx, space_id, index_name).expect("load should succeed");
    assert!(
        after.is_none(),
        "CatchingUp state should be discarded on crash recovery"
    );
}

#[test]
fn test_crash_recovery_preserves_publishing_state() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "publishing_crash_idx";

    // Simulate crash during Publishing
    let mut state = generation_state(1, 10);
    state
        .transition_to_catching_up()
        .expect("Building should transition to CatchingUp");
    state
        .transition_to_publishing(CommitLsn::new(50))
        .expect("CatchingUp should transition to Publishing");
    save_generation_build_state(&ctx, space_id, index_name, &state)
        .expect("Publishing state should save");

    // Recover (should NOT discard Publishing state since the manifest may be published)
    resolve_crash_recovery(&ctx, space_id, index_name).expect("recovery should succeed");

    let after =
        load_generation_build_state(&ctx, space_id, index_name).expect("load should succeed");
    assert!(
        after.is_some(),
        "Publishing state should be preserved on crash recovery"
    );
    assert_eq!(after.unwrap().state, GenerationState::Publishing);
}

#[test]
fn test_generation_build_state_not_found_for_missing_index() {
    let (_temp, ctx) = setup_persistent_context();
    let space_id = 1u64;
    let index_name = "nonexistent_idx";

    let result =
        load_generation_build_state(&ctx, space_id, index_name).expect("load should succeed");
    assert!(result.is_none());
}

#[test]
fn test_remove_generation_build_state_on_clean_state_is_noop() {
    let (_temp, ctx) = setup_persistent_context();
    let result = remove_generation_build_state(&ctx, 1u64, "no_state_idx");
    assert!(result.is_ok());
}

#[test]
fn test_generation_build_state_transitions() {
    let building = generation_state(1, 10);
    assert!(!building.is_active());

    let mut active = generation_state(1, 10);
    active
        .transition_to_catching_up()
        .expect("Building should transition to CatchingUp");
    active
        .transition_to_publishing(CommitLsn::new(10))
        .expect("CatchingUp should transition to Publishing");
    active
        .transition_to_active()
        .expect("Publishing should transition to Active");
    assert!(active.is_active());
}

#[test]
fn test_build_vertex_index_data_handles_empty_input() {
    let index = Index::new(IndexConfig {
        id: 1,
        name: "person_name_idx".to_string(),
        space_id: 0,
        schema_name: "Person".to_string(),
        fields: vec![IndexField::new(
            "name".to_string(),
            Value::string(""),
            false,
        )],
        properties: vec![],
        index_type: IndexType::TagIndex,
        is_unique: false,
        covering: false,
        partial_condition: None,
    });

    let vertices = vec![];
    let (forward, reverse) =
        build_vertex_index_data(0, &index, &vertices, linkrs_core::types::MAX_TIMESTAMP)
            .expect("build should succeed with empty input");
    assert!(forward.is_empty(), "forward map should be empty");
    assert!(reverse.is_empty(), "reverse map should be empty");
}

#[test]
fn test_flush_index_data_writes_valid_files() {
    use crate::index::types::IndexRecord;
    use linkrs_core::types::MAX_TIMESTAMP;
    use std::collections::BTreeMap;
    use std::fs;

    let temp_dir = tempfile::TempDir::new().expect("Failed to create temp dir");
    let gen_dir = temp_dir.path().join("gen_1");
    fs::create_dir_all(&gen_dir).expect("dir should be created");

    let mut forward = BTreeMap::new();
    forward.insert(vec![0, 1, 2], IndexRecord::new(MAX_TIMESTAMP));
    let mut reverse = BTreeMap::new();
    reverse.insert(vec![3, 4, 5], IndexRecord::new(MAX_TIMESTAMP));

    write_generation_checkpoint(&gen_dir, &forward, &reverse).expect("flush should succeed");

    assert!(
        gen_dir.join("forward_chunks/chunk_index.bin").exists(),
        "forward chunk index should exist"
    );
    assert!(
        gen_dir.join("reverse_chunks/chunk_index.bin").exists(),
        "reverse chunk index should exist"
    );
}