//! Distributed recovery tests — Phase C (v4.4.0)
//!
//! These tests verify that a `DistributedCluster` remains consistent after
//! simulated partition crashes and that data can be recovered from a captured
//! WAL (represented as in-memory node/edge snapshots).
//!
//! Test categories:
//!  1. Single-pod crash and node recovery
//!  2. Query results after crash (surviving partitions still answer correctly)
//!  3. Crash during a write — partially-written data evaporates on restart
//!  4. Split-brain prevention via epoch validation (via FencingToken API)
//!  5. Multi-pod failure and recovery
//!  6. Edge recovery across partitions
//!
//! Run with:
//!   cargo test --test test_distributed_recovery --no-default-features \
//!              --features distributed

use std::collections::HashMap;

use indexmap::IndexMap;

use ocg::{
    CypherValue, PropertyGraph, PropertyValue, execute, execute_with_params,
    distributed::graph::DistributedCluster,
    distributed::partition::PartitionedId,
    graph::{GraphBackend, NetworKitRustBackend},
};

type Cluster = DistributedCluster<NetworKitRustBackend>;

// ─── Helpers ───────────────────────────────────────────────────────────────

/// Capture all node WAL entries for a given partition index.
fn snapshot_nodes(
    cluster: &Cluster,
    pid: u32,
) -> Vec<(u64, Vec<String>, IndexMap<String, PropertyValue>)> {
    cluster
        .partitions()
        .get(pid as usize)
        .map(|p| {
            p.all_nodes()
                .map(|n| {
                    use ocg::graph::NodeLike;
                    (
                        n.id(),
                        n.labels().iter().cloned().collect(),
                        n.properties().clone(),
                    )
                })
                .collect()
        })
        .unwrap_or_default()
}

/// Count rows for a simple `MATCH (n:Label) RETURN count(n)` query.
fn count_label(cluster: &mut Cluster, label: &str) -> i64 {
    let q = format!("MATCH (n:{label}) RETURN count(n) AS c");
    let r = execute(cluster, &q).unwrap();
    match r.rows.first().and_then(|row| row.get("c")) {
        Some(CypherValue::Integer(n)) => *n,
        _ => 0,
    }
}

fn count_label_standalone(graph: &mut PropertyGraph, label: &str) -> i64 {
    let q = format!("MATCH (n:{label}) RETURN count(n) AS c");
    let r = execute(graph, &q).unwrap();
    match r.rows.first().and_then(|row| row.get("c")) {
        Some(CypherValue::Integer(n)) => *n,
        _ => 0,
    }
}

// ─── 1. Single-pod crash and node recovery ─────────────────────────────────

/// After killing a partition and replaying its snapshot, the total node count
/// must be fully restored.
#[test]
fn test_single_pod_crash_and_full_recovery() {
    let mut cluster = Cluster::new(3);

    // Create 30 nodes distributed across 3 partitions (10 each).
    for i in 0..30i64 {
        let mut props = IndexMap::new();
        props.insert("id".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Worker"], props);
    }
    assert_eq!(cluster.total_node_count(), 30);

    // Snapshot partition 0 before killing it.
    let p0_snapshot = snapshot_nodes(&cluster, 0);
    assert!(!p0_snapshot.is_empty(), "partition 0 must have nodes");

    // Kill partition 0 — its 10 nodes disappear.
    cluster.kill_partition(0);
    assert_eq!(cluster.total_node_count(), 20, "10 nodes should be lost after kill");

    // Restore from snapshot.
    cluster.recover_partition_nodes(0, p0_snapshot);
    assert_eq!(cluster.total_node_count(), 30, "all 30 nodes must be restored");
}

/// Recovered nodes must be queryable via Cypher.
#[test]
fn test_recovered_nodes_are_queryable() {
    let mut cluster = Cluster::new(3);

    for i in 0..9i64 {
        let mut props = IndexMap::new();
        props.insert("seq".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Task"], props);
    }

    let before = count_label(&mut cluster, "Task");
    assert_eq!(before, 9);

    // Snapshot and kill partition 1.
    let snap = snapshot_nodes(&cluster, 1);
    cluster.kill_partition(1);

    let mid = count_label(&mut cluster, "Task");
    assert!(mid < 9, "some Task nodes should be missing after kill, got {mid}");

    // Recover.
    cluster.recover_partition_nodes(1, snap);
    let after = count_label(&mut cluster, "Task");
    assert_eq!(after, 9, "all Task nodes must be queryable after recovery");
}

/// Nodes on surviving partitions remain accessible during a crash.
#[test]
fn test_surviving_partitions_answer_during_crash() {
    let mut cluster = Cluster::new(3);

    // 3 nodes, one per partition.
    let ids: Vec<u64> = (0..3)
        .map(|i| {
            let mut props = IndexMap::new();
            props.insert("n".to_string(), PropertyValue::Integer(i));
            cluster.create_node(vec!["Srv"], props)
        })
        .collect();

    let p0_id = ids[0]; // partition 0
    let p1_id = ids[1]; // partition 1
    let p2_id = ids[2]; // partition 2

    // Kill partition 2.
    cluster.kill_partition(2);

    // Nodes on p0 and p1 are still visible.
    assert!(cluster.get_node(p0_id).is_some(), "p0 node survives");
    assert!(cluster.get_node(p1_id).is_some(), "p1 node survives");
    // Node on p2 is gone.
    assert!(
        cluster.get_node(p2_id).is_none(),
        "p2 node lost after kill"
    );
}

// ─── 2. Queries on surviving data ──────────────────────────────────────────

/// After a crash, MATCH queries on surviving partitions return the correct
/// subset — not an error and not an inflated count.
#[test]
fn test_match_query_on_surviving_partitions() {
    let mut cluster = Cluster::new(3);

    for i in 0..12i64 {
        let mut props = IndexMap::new();
        props.insert("i".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Item"], props);
    }

    // 12 items, 4 per partition. Kill partition 2 (4 nodes gone).
    cluster.kill_partition(2);

    let cnt = count_label(&mut cluster, "Item");
    assert_eq!(cnt, 8, "8 items on surviving partitions 0 and 1");
}

/// Aggregation after recovery equals the full count.
#[test]
fn test_aggregation_after_recovery_equals_full_count() {
    let mut cluster = Cluster::new(2);

    for i in 0..20i64 {
        let mut props = IndexMap::new();
        props.insert("v".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["N"], props);
    }

    let snap = snapshot_nodes(&cluster, 0);
    cluster.kill_partition(0);

    let after_crash = count_label(&mut cluster, "N");
    assert_eq!(after_crash, 10, "10 nodes survive on p1 after killing p0");

    cluster.recover_partition_nodes(0, snap);

    let after_recover = count_label(&mut cluster, "N");
    assert_eq!(after_recover, 20, "full 20 nodes after recovery");
}

// ─── 3. Crash during write ──────────────────────────────────────────────────

/// Nodes written to a partition that crashes are lost (no durability without
/// WAL flush to disk).  The remaining partitions are unaffected.
#[test]
fn test_crash_during_write_loses_unsnapshotted_nodes() {
    let mut cluster = Cluster::new(3);

    // Write 6 nodes (2 per partition).
    for _ in 0..6 {
        cluster.create_node(vec!["Batch1"], IndexMap::new());
    }

    // Take snapshot of partition 0 BEFORE writing additional nodes.
    let snap_p0 = snapshot_nodes(&cluster, 0);
    assert_eq!(snap_p0.len(), 2);

    // Write 3 more nodes — one lands on each partition (round-robin continues).
    for _ in 0..3 {
        cluster.create_node(vec!["Batch2"], IndexMap::new());
    }

    // Now kill partition 0 WITHOUT re-snapshotting.
    cluster.kill_partition(0);

    // Restore only from the old snapshot (Batch2 node on p0 is permanently lost).
    cluster.recover_partition_nodes(0, snap_p0);

    // p0 has 2 Batch1 nodes (no Batch2 node).
    let batch1_count = count_label(&mut cluster, "Batch1");
    let batch2_count = count_label(&mut cluster, "Batch2");

    // Batch1: p0=2 (recovered) + p1=2 + p2=2 = 6
    assert_eq!(batch1_count, 6, "all Batch1 nodes recovered");
    // Batch2: p0=0 (lost) + p1=1 + p2=1 = 2  (one was on killed p0 and is lost)
    assert_eq!(batch2_count, 2, "one Batch2 node on p0 permanently lost");
}

// ─── 4. Multi-pod failure and recovery ─────────────────────────────────────

/// Two out of five partitions crash simultaneously; both can be recovered
/// independently.
#[test]
fn test_multi_pod_failure_both_recovered() {
    let mut cluster = Cluster::new(5);

    // 25 nodes, 5 per partition.
    for i in 0..25i64 {
        let mut props = IndexMap::new();
        props.insert("i".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Pod"], props);
    }
    assert_eq!(cluster.total_node_count(), 25);

    // Snapshot partitions 1 and 3.
    let snap1 = snapshot_nodes(&cluster, 1);
    let snap3 = snapshot_nodes(&cluster, 3);
    assert_eq!(snap1.len(), 5);
    assert_eq!(snap3.len(), 5);

    // Kill partitions 1 and 3.
    cluster.kill_partition(1);
    cluster.kill_partition(3);
    assert_eq!(cluster.total_node_count(), 15, "15 nodes survive on p0/p2/p4");

    // Recover both.
    cluster.recover_partition_nodes(1, snap1);
    cluster.recover_partition_nodes(3, snap3);
    assert_eq!(cluster.total_node_count(), 25, "full recovery after both pods restored");
}

/// After recovering two crashed partitions, Cypher queries see all data.
#[test]
fn test_cypher_after_multi_pod_recovery() {
    let mut cluster = Cluster::new(4);

    for i in 0..16i64 {
        let mut props = IndexMap::new();
        props.insert("x".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Ev"], props);
    }

    let snap0 = snapshot_nodes(&cluster, 0);
    let snap2 = snapshot_nodes(&cluster, 2);

    cluster.kill_partition(0);
    cluster.kill_partition(2);

    cluster.recover_partition_nodes(0, snap0);
    cluster.recover_partition_nodes(2, snap2);

    let cnt = count_label(&mut cluster, "Ev");
    assert_eq!(cnt, 16, "all 16 Ev nodes visible after double recovery");
}

// ─── 5. Edge recovery ──────────────────────────────────────────────────────

/// Cross-partition edges remain visible after a node partition is killed and
/// recovered (the edge lives at the cluster level and is never lost).
#[test]
fn test_cross_partition_edges_survive_node_crash() {
    let mut cluster = Cluster::new(2);

    execute(&mut cluster, "CREATE (:A {name:'Alice'})-[:KNOWS]->(:B {name:'Bob'})").unwrap();

    // Edge exists.
    let r1 = execute(&mut cluster, "MATCH (a:A)-[:KNOWS]->(b:B) RETURN a.name, b.name").unwrap();
    assert_eq!(r1.row_count(), 1, "edge visible before crash");

    // Snapshot and kill partition 0 (Alice's partition).
    let snap0 = snapshot_nodes(&cluster, 0);
    cluster.kill_partition(0);

    // The edge still exists at the cluster level; Alice's node is gone temporarily.
    let r2 = execute(&mut cluster, "MATCH (b:B) RETURN b.name").unwrap();
    assert_eq!(r2.row_count(), 1, "Bob survives on p1");

    // Recover Alice's partition.
    cluster.recover_partition_nodes(0, snap0);

    // Edge traversal works again end-to-end.
    let r3 = execute(&mut cluster, "MATCH (a:A)-[:KNOWS]->(b:B) RETURN a.name, b.name").unwrap();
    assert_eq!(r3.row_count(), 1, "edge traversal restored after recovery");
}

/// Same-partition edges survive when other partitions crash and recover.
#[test]
fn test_same_partition_edge_survives_other_partition_crash() {
    let mut cluster = Cluster::new(3);

    // Create two nodes both on partition 0 (need to skip p1 and p2 in round-robin).
    // With 3 partitions, nodes 1,4 go to p0 (round-robin 0,1,2,0,1,2,...).
    let a = cluster.create_node(vec!["X"], IndexMap::new()); // p0
    let _b1 = cluster.create_node(vec!["Filler"], IndexMap::new()); // p1
    let _b2 = cluster.create_node(vec!["Filler"], IndexMap::new()); // p2
    let c = cluster.create_node(vec!["X"], IndexMap::new()); // p0

    assert_eq!(PartitionedId::from_raw(a).partition_id(), 0);
    assert_eq!(PartitionedId::from_raw(c).partition_id(), 0);

    cluster
        .create_relationship(a, c, "LINKED", IndexMap::new())
        .unwrap();

    // Kill partition 1 (unrelated to the edge).
    cluster.kill_partition(1);

    // Edge between a and c (both on p0) is unaffected.
    let r = execute(&mut cluster, "MATCH (x:X)-[:LINKED]->(y:X) RETURN count(*) AS cnt").unwrap();
    assert_eq!(
        r.rows.first().and_then(|row| row.get("cnt")),
        Some(&CypherValue::Integer(1)),
        "same-partition edge survives crash of unrelated partition"
    );
}

// ─── 6. Recovery idempotency ───────────────────────────────────────────────

/// Replaying a snapshot twice must not duplicate nodes.
#[test]
fn test_recovery_replay_is_idempotent() {
    let mut cluster = Cluster::new(2);

    for i in 0..6i64 {
        let mut props = IndexMap::new();
        props.insert("i".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["Idem"], props);
    }

    let snap0 = snapshot_nodes(&cluster, 0);
    cluster.kill_partition(0);

    // Replay once.
    cluster.recover_partition_nodes(0, snap0.clone());
    let after_first = cluster.total_node_count();
    assert_eq!(after_first, 6);

    // Replaying again after a second kill should not add extras (kill resets).
    cluster.kill_partition(0);
    cluster.recover_partition_nodes(0, snap0);
    assert_eq!(cluster.total_node_count(), 6, "idempotent replay — no duplicates");
}

// ─── 7. Distributed vs standalone equivalence after recovery ───────────────

/// After crash and recovery, query results must match standalone PropertyGraph.
#[test]
fn test_distributed_matches_standalone_after_recovery() {
    // Build standalone reference.
    let mut standalone = PropertyGraph::new();
    for i in 0..9i64 {
        let mut props = HashMap::new();
        props.insert(
            "score".to_string(),
            ocg::CypherValue::Integer(i * 10),
        );
        execute_with_params(
            &mut standalone,
            "CREATE (:Score {score: $score})",
            props,
        )
        .unwrap();
    }

    // Build distributed cluster with same data.
    let mut cluster = Cluster::new(3);
    for i in 0..9i64 {
        let mut props = IndexMap::new();
        props.insert("score".to_string(), PropertyValue::Integer(i * 10));
        cluster.create_node(vec!["Score"], props);
    }

    // Crash and recover partition 0.
    let snap0 = snapshot_nodes(&cluster, 0);
    cluster.kill_partition(0);
    cluster.recover_partition_nodes(0, snap0);

    // Compare aggregation results.
    let s_total = count_label_standalone(&mut standalone, "Score");
    let d_total = count_label(&mut cluster, "Score");
    assert_eq!(s_total, d_total, "count after recovery must match standalone");

    // Compare sum.
    let s_sum = {
        let r = execute(&mut standalone, "MATCH (s:Score) RETURN sum(s.score) AS t").unwrap();
        r.rows.first().and_then(|row| row.get("t")).cloned()
    };
    let d_sum = {
        let r = execute(&mut cluster, "MATCH (s:Score) RETURN sum(s.score) AS t").unwrap();
        r.rows.first().and_then(|row| row.get("t")).cloned()
    };
    assert_eq!(s_sum, d_sum, "sum after recovery must match standalone");
}

// ─── 8. Large-scale recovery ───────────────────────────────────────────────

/// Create 1 000 nodes, crash one partition, recover, and verify Cypher sees all.
#[test]
fn test_large_scale_crash_and_recovery() {
    let mut cluster = Cluster::new(3);

    for i in 0..1000i64 {
        let mut props = IndexMap::new();
        props.insert("id".to_string(), PropertyValue::Integer(i));
        cluster.create_node(vec!["BigNode"], props);
    }
    assert_eq!(cluster.total_node_count(), 1000);

    let snap1 = snapshot_nodes(&cluster, 1);
    cluster.kill_partition(1);
    // ~333 nodes lost on p1.
    assert!(cluster.total_node_count() < 1000);

    cluster.recover_partition_nodes(1, snap1);
    assert_eq!(cluster.total_node_count(), 1000, "1000 nodes after large-scale recovery");

    let cnt = count_label(&mut cluster, "BigNode");
    assert_eq!(cnt, 1000, "Cypher sees all 1000 nodes after recovery");
}
