use std::collections::{HashMap, HashSet};

use assertables::{assume, assume_eq};
use decorum::R64;
use indexmap::IndexMap;
use rand::SeedableRng;
use rand_pcg::Pcg64;

use anysystem::logger::LogEntry;
use anysystem::test::TestResult;
use anysystem::{Message, System};

use crate::common::*;

pub fn test_single_node(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    let proc = "0";
    let key = random_string(8, &mut rand).to_uppercase();
    let value = random_string(8, &mut rand);
    let max_steps = 10;

    check_get(&mut sys, proc, &key, None, max_steps)?;
    check_put(&mut sys, proc, &key, &value, max_steps)?;
    check_get(&mut sys, proc, &key, Some(&value), max_steps)?;
    let updated = format!("{value}-updated");
    check_put(&mut sys, proc, &key, &updated, max_steps)?;
    check_get(&mut sys, proc, &key, Some(&updated), max_steps)?;
    check_delete(&mut sys, proc, &key, Some(&updated), max_steps)?;
    check_get(&mut sys, proc, &key, None, max_steps)?;
    check_delete(&mut sys, proc, &key, None, max_steps)
}

pub fn test_inserts(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    // insert random key-value pairs from each node
    let mut kv = IndexMap::new();
    for proc in sys.process_names() {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // check that all key-values can be read from each node
    let procs = sys.process_names();
    check(&mut sys, &procs, &kv, true, false)?;
    // Update through a different entry node and read the new values everywhere.
    for (i, (key, value)) in kv.iter_mut().enumerate() {
        value.push_str("-updated");
        check_put(&mut sys, &procs[(i + 1) % procs.len()], key, value, 100)?;
    }
    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_deletes(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut kv = IndexMap::new();

    // insert random key-value pairs from each node
    for proc in sys.process_names() {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // delete each key from one node and check that key is not present from another
    for (k, v) in kv.iter() {
        let read_proc = random_proc(&sys, &mut rand);
        let mut delete_proc = random_proc(&sys, &mut rand);
        while delete_proc == read_proc {
            delete_proc = random_proc(&sys, &mut rand);
        }
        check_get(&mut sys, &read_proc, k, Some(v), 100)?;
        check_delete(&mut sys, &delete_proc, k, Some(v), 100)?;
        check_get(&mut sys, &read_proc, k, None, 100)?;
    }

    kv.clear();
    let procs = sys.process_names();
    check(&mut sys, &procs, &kv, false, false)
}

const MEMORY_NEW_KEYS_PER_NODE: usize = 1000;

pub fn test_memory_growth(config: &TestConfig) -> TestResult {
    let small = memory_growth(config, 5)?;
    let large = memory_growth(config, 20)?;
    check_memory_growth(small, large)
}

fn check_memory_growth(bytes_at_5: u64, bytes_at_20: u64) -> TestResult {
    // Compare equal numbers of new keys before rounding either measurement.
    let extra = i128::from(bytes_at_20) - 4 * i128::from(bytes_at_5);
    let new_keys = (20 * MEMORY_NEW_KEYS_PER_NODE) as i128;
    let difference = extra as f64 / new_keys as f64;
    println!("Additional bytes per new key (20 vs 5 nodes): {difference:.5}");
    assume!(
        extra <= 64 * new_keys,
        format!("Memory growth exceeds 64 extra bytes per new key ({difference:.5} observed)")
    )
}

fn memory_growth(config: &TestConfig, nodes: u32) -> Result<u64, String> {
    let mut sys = build_system(
        &TestConfig {
            proc_count: nodes,
            ..*config
        },
        true,
    );
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut keys = HashSet::new();
    let mut measure_at = |records: usize| -> Result<u64, String> {
        while keys.len() < records {
            let key = random_string(8, &mut rand).to_uppercase();
            if !keys.insert(key.clone()) {
                continue;
            }
            let value = random_string(8, &mut rand);
            let ingress = random_proc(&sys, &mut rand);
            check_put(&mut sys, &ingress, &key, &value, 100)?;
        }
        // Allow pending work to finish without waiting forever for recurring timers.
        sys.steps(1000);
        let mut stored = 0;
        let mut bytes = 0;
        for proc in sys.process_names() {
            stored += count_records(&mut sys, &proc)?;
            bytes += sys.max_size(&proc);
        }
        assume_eq!(stored, records as u64)?;
        Ok(bytes)
    };

    let before = 500 * nodes as usize;
    let after = before + MEMORY_NEW_KEYS_PER_NODE * nodes as usize;
    let initial = measure_at(before)?;
    let final_size = measure_at(after)?;
    let bytes = final_size - initial;
    let growth = bytes as f64 / (after - before) as f64;
    println!("{nodes} nodes, {before} to {after} records: {growth:.3} bytes per new key");
    Ok(bytes)
}

pub fn test_zero_hop(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let procs = sys.process_names();
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut kv = IndexMap::new();
    for _ in 0..16 {
        let key = random_string(8, &mut rand).to_uppercase();
        let value = random_string(8, &mut rand);
        let ingress = random_proc(&sys, &mut rand);
        check_put(&mut sys, &ingress, &key, &value, 100)?;
        kv.insert(key, value);
    }
    sys.steps(1000);
    let mut owners = HashMap::new();
    for proc in &procs {
        for key in dump_keys(&mut sys, proc)? {
            assume!(
                owners.insert(key, proc).is_none(),
                "Key stored on multiple nodes"
            )?;
        }
    }
    for (key, value) in kv {
        let owner = owners
            .get(&key)
            .ok_or_else(|| format!("No owner for {key}"))?;
        for ingress in &procs {
            let start = sys.logger().trace().len();
            check_get(&mut sys, ingress, &key, Some(&value), 100)?;
            check_read_route(&sys.logger().trace()[start..], ingress, owner)
                .map_err(|error| format!("GET({key}) via {ingress}, owner {owner}: {error}"))?;
        }
    }
    Ok(true)
}

fn check_read_route(trace: &[LogEntry], ingress: &str, owner: &str) -> TestResult {
    // Follow the GET's messages and timers so unrelated background traffic
    // interleaved with the read does not count as a routing hop.
    let mut messages = HashSet::new();
    let mut timers = HashSet::new();
    let mut handling_read = true;
    let mut reached_owner = ingress == owner;
    let mut received_reply = ingress == owner;
    for entry in trace {
        match entry {
            LogEntry::MessageSent {
                msg_id,
                src_proc,
                dst_proc,
                ..
            } if handling_read => {
                assume!(
                    dst_proc == ingress || dst_proc == owner,
                    format!("Indirect route through {dst_proc} (sent by {src_proc})")
                )?;
                messages.insert(msg_id);
            }
            LogEntry::MessageReceived {
                msg_id,
                src_proc,
                dst_proc,
                ..
            } => {
                handling_read = messages.contains(msg_id);
                // A process may also dispatch queued requests from an existing timer.
                reached_owner |= src_proc == ingress && dst_proc == owner;
                received_reply |= reached_owner && src_proc == owner && dst_proc == ingress;
            }
            LogEntry::TimerSet { timer_id, .. } if handling_read => {
                timers.insert(timer_id);
            }
            LogEntry::TimerFired { timer_id, .. } => handling_read = timers.contains(timer_id),
            LogEntry::LocalMessageSent { proc, .. } if proc == ingress => break,
            _ => {}
        }
    }
    assume!(
        received_reply,
        "No direct exchange with the owner before GET_RESP".to_owned()
    )
}

pub fn test_node_added(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    // insert random key-value pairs
    let keys_count = 100;
    let mut kv = IndexMap::new();
    for _ in 0..keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // add new node to the system
    let added = sys.process_names().len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);

    // run the system until key the distribution is stabilized
    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;

    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_node_added_during_inserts(config: &TestConfig) -> TestResult {
    for (scenario, delays, update_existing) in [
        ("random delays", None, false),
        ("delayed PUT delivery", Some((0.1, 0.01)), false),
        ("updates with delayed migration", Some((0.01, 0.1)), true),
        ("updates with delayed PUT delivery", Some((0.1, 0.01)), true),
    ] {
        println!("Scenario: {scenario} (seed {})", config.seed);
        node_added_during_inserts_scenario(config, delays, update_existing)
            .map_err(|error| format!("{scenario}: {error}"))?;
    }
    for (scenario, delays) in [
        ("new PUTs with random delays", None),
        ("new PUTs with delayed migration", Some((0.1, 0.01))),
        ("new PUTs with delayed PUT delivery", Some((0.01, 0.1))),
    ] {
        for via_new_node in [true, false] {
            let entry = if via_new_node {
                "new node"
            } else {
                "old nodes"
            };
            println!("Scenario: {scenario} (seed {}, via {entry})", config.seed);
            node_added_during_new_puts_scenario(config, delays, via_new_node)
                .map_err(|error| format!("{scenario}: {error}"))?;
        }
    }
    Ok(true)
}

fn node_added_during_inserts_scenario(
    config: &TestConfig,
    delays: Option<(f64, f64)>,
    update_existing: bool,
) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    // Prepare unique keys and, for update scenarios, finish the initial writes.
    let keys_count = 100;
    let mut kv = IndexMap::new();
    let mut requests = Vec::new();
    while kv.len() < keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        if kv.contains_key(&k) {
            continue;
        }
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        if update_existing {
            check_put(&mut sys, &proc, &k, &format!("old-{v}"), 100)?;
        }
        requests.push((proc, k.clone(), v.clone()));
        kv.insert(k, v);
    }
    if let Some((put_delay, _)) = delays {
        sys.network().set_delays(put_delay, put_delay);
    }
    let mut expected: IndexMap<String, IndexMap<String, String>> = IndexMap::new();
    for (proc, key, value) in requests {
        sys.send_local_message(
            &proc,
            Message::json(
                "PUT",
                &PutReqMessage {
                    key: &key,
                    value: &value,
                },
            ),
        );
        expected.entry(proc).or_default().insert(key, value);
    }

    // Already-sent PUTs keep their delivery times. Test both orders relative
    // to migration, within the usual 0.01..0.1 delay range.
    if let Some((_, migration_delay)) = delays {
        sys.network().set_delays(migration_delay, migration_delay);
    }

    // add new node to the system (while puts are processed)
    let added = sys.process_names().len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);
    if delays.is_some() {
        // Forwarded PUTs can reach the new owner before a delayed snapshot.
        sys.network().set_delays(0.01, 0.01);
    }

    // wait for completion of puts
    for (proc, responses) in expected {
        wait_put_responses(&mut sys, &proc, responses, 1000)?;
    }

    // run the system until key the distribution is stabilized
    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;

    check(&mut sys, &procs, &kv, true, false)
}

fn node_added_during_new_puts_scenario(
    config: &TestConfig,
    delays: Option<(f64, f64)>,
    via_new_node: bool,
) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut kv = IndexMap::new();
    while kv.len() < 100 {
        let key = random_string(8, &mut rand).to_uppercase();
        if kv.contains_key(&key) {
            continue;
        }
        let value = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &key, &value, 100)?;
        kv.insert(key, value);
    }
    let old_procs = sys.process_names();
    step_until_stabilized(&mut sys, &old_procs, kv.len() as u64, 100, 1000)?;

    // Update half the existing records and leave the rest unchanged.
    // Each key has only one concurrent PUT, so the final value is unambiguous.
    let mut requests = Vec::new();
    for (key, value) in kv.iter_mut().step_by(2) {
        value.push_str("-updated");
        requests.push((key.clone(), value.clone()));
    }
    // Also insert new keys while the existing records are being moved.
    while requests.len() < 100 {
        let key = random_string(8, &mut rand).to_uppercase();
        if kv.contains_key(&key) {
            continue;
        }
        let value = random_string(8, &mut rand);
        requests.push((key.clone(), value.clone()));
        kv.insert(key, value);
    }

    if let Some((migration_delay, _)) = delays {
        sys.network().set_delays(migration_delay, migration_delay);
    }
    let added = old_procs.len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);
    // Do not advance the simulation: new PUTs race with pending transfers.
    if let Some((_, put_delay)) = delays {
        sys.network().set_delays(put_delay, put_delay);
    }
    let mut expected: IndexMap<String, IndexMap<String, String>> = IndexMap::new();
    for (i, (key, value)) in requests.into_iter().enumerate() {
        // Run the same workload through both entry paths in separate systems.
        let proc = if via_new_node {
            added.clone()
        } else {
            old_procs[i % old_procs.len()].clone()
        };
        sys.send_local_message(
            &proc,
            Message::json(
                "PUT",
                &PutReqMessage {
                    key: &key,
                    value: &value,
                },
            ),
        );
        expected.entry(proc).or_default().insert(key, value);
    }
    for (proc, responses) in expected {
        wait_put_responses(&mut sys, &proc, responses, 1000)?;
    }

    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_node_added_during_deletions(config: &TestConfig) -> TestResult {
    for (scenario, delays) in [
        ("random delays", None),
        ("delayed migration", Some((0.1, 0.01))),
        ("delayed DELETE delivery", Some((0.01, 0.1))),
    ] {
        println!("Scenario: {scenario} (seed {})", config.seed);
        node_added_during_deletions_scenario(config, delays)
            .map_err(|error| format!("{scenario}: {error}"))?;
    }
    Ok(true)
}

fn node_added_during_deletions_scenario(
    config: &TestConfig,
    delays: Option<(f64, f64)>,
) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut kv = IndexMap::new();
    for _ in 0..100 {
        let key = random_string(8, &mut rand).to_uppercase();
        let value = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &key, &value, 100)?;
        kv.insert(key, value);
    }
    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;

    // Keep some records to detect solutions that discard the entire transfer.
    let deleted: Vec<String> = kv.keys().step_by(2).cloned().collect();
    let mut requests: IndexMap<String, Option<String>> = deleted
        .iter()
        .map(|key| (key.clone(), kv.shift_remove(key)))
        .collect();
    for i in 0..5 {
        requests.insert(format!("missing-{i}"), None);
    }

    if let Some((migration_delay, _)) = delays {
        sys.network().set_delays(migration_delay, migration_delay);
    }
    let added = sys.process_names().len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);
    // Do not advance the simulation: deletions race with the pending transfers.
    if let Some((_, delete_delay)) = delays {
        sys.network().set_delays(delete_delay, delete_delay);
    }
    let mut expected: IndexMap<String, IndexMap<String, Option<String>>> = IndexMap::new();
    for (i, (key, value)) in requests.iter().enumerate() {
        let proc = if i % 2 == 0 {
            added.clone()
        } else {
            random_proc(&sys, &mut rand)
        };
        sys.send_local_message(&proc, Message::json("DELETE", &DeleteReqMessage { key }));
        expected
            .entry(proc)
            .or_default()
            .insert(key.clone(), value.clone());
    }
    for (proc, responses) in expected {
        wait_delete_responses(&mut sys, &proc, responses, 1000)?;
    }

    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    check(&mut sys, &procs, &kv, true, false)?;
    for proc in &procs {
        for key in requests.keys() {
            check_get(&mut sys, proc, key, None, 100)?;
        }
    }
    // Deleted keys must remain usable after migration (no permanent tombstones).
    for key in deleted {
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &key, "reinserted", 100)?;
        kv.insert(key, "reinserted".to_string());
    }
    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_node_added_during_reads(config: &TestConfig) -> TestResult {
    for (scenario, delays) in [
        ("random delays", None),
        ("delayed migration", Some((0.1, 0.01))),
        ("delayed GET delivery", Some((0.01, 0.1))),
    ] {
        for via_new_node in [true, false] {
            let entry = if via_new_node {
                "new node"
            } else {
                "old nodes"
            };
            println!("Scenario: {scenario} (seed {}, via {entry})", config.seed);
            node_added_during_reads_scenario(config, delays, via_new_node)
                .map_err(|error| format!("{scenario} via {entry}: {error}"))?;
        }
    }
    Ok(true)
}

fn node_added_during_reads_scenario(
    config: &TestConfig,
    delays: Option<(f64, f64)>,
    via_new_node: bool,
) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut kv = IndexMap::new();
    while kv.len() < 100 {
        let key = random_string(8, &mut rand).to_uppercase();
        if kv.contains_key(&key) {
            continue;
        }
        let value = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &key, &value, 100)?;
        kv.insert(key, value);
    }
    let old_procs = sys.process_names();
    step_until_stabilized(&mut sys, &old_procs, kv.len() as u64, 100, 1000)?;

    let mut requests: IndexMap<String, Option<String>> = kv
        .iter()
        .map(|(key, value)| (key.clone(), Some(value.clone())))
        .collect();
    for i in 0..5 {
        requests.insert(format!("missing-{i}"), None);
    }

    if let Some((migration_delay, _)) = delays {
        sys.network().set_delays(migration_delay, migration_delay);
    }
    let added = old_procs.len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);
    // Send every GET before advancing the simulation; no concurrent writes.
    if let Some((_, get_delay)) = delays {
        sys.network().set_delays(get_delay, get_delay);
    }
    let mut expected: IndexMap<String, IndexMap<String, Option<String>>> = IndexMap::new();
    for (i, (key, value)) in requests.into_iter().enumerate() {
        let proc = if via_new_node {
            added.clone()
        } else {
            old_procs[i % old_procs.len()].clone()
        };
        sys.send_local_message(&proc, Message::json("GET", &GetReqMessage { key: &key }));
        expected.entry(proc).or_default().insert(key, value);
    }
    for (proc, responses) in expected {
        wait_get_responses(&mut sys, &proc, responses, 1000)?;
    }

    // Reading must leave both values and the set of stored keys unchanged.
    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_node_removed(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    // insert random key-value pairs
    let keys_count = 100;
    let mut kv = IndexMap::new();
    for _ in 0..keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // remove a node from the system
    let removed = random_proc(&sys, &mut rand);
    let count = count_records(&mut sys, &removed)?;
    assume!(count > 0, "Node stores no records, bad distribution")?;
    send_node_removed(&mut sys, &removed);

    // run the system until key the distribution is stabilized
    let procs: Vec<String> = sys
        .process_names()
        .into_iter()
        .filter(|x| *x != removed)
        .collect();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;

    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_node_removed_after_crash(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);

    // insert random key-value pairs
    let keys_count = 100;
    let mut kv = IndexMap::new();
    for _ in 0..keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // crash a node and remove it from the system (stored keys are lost)
    let crashed = random_proc(&sys, &mut rand);
    let crashed_keys = dump_keys(&mut sys, &crashed)?;
    assume!(
        !crashed_keys.is_empty(),
        "Node stores no records, bad distribution"
    )?;
    for k in crashed_keys {
        kv.shift_remove(&k);
    }
    sys.crash_node(&sys.proc_node_name(&crashed));
    send_node_removed(&mut sys, &crashed);

    // run the system until key the distribution is stabilized
    let procs: Vec<String> = sys
        .process_names()
        .into_iter()
        .filter(|x| *x != crashed)
        .collect();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;

    check(&mut sys, &procs, &kv, true, false)
}

pub fn test_migration(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut procs = sys.process_names();

    // insert random key-value pairs
    let keys_count = 10000;
    let mut kv = IndexMap::new();
    for _ in 0..keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // add new N nodes to the system
    for i in 0..config.proc_count {
        let added = format!("{}", config.proc_count + i);
        add_node(&added, &mut sys, config);
        send_node_added(&mut sys, &added);
        procs.push(added);
        step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    }

    check(&mut sys, &procs, &kv, false, false)?;

    // remove old N nodes
    for _ in 0..config.proc_count {
        let removed = &procs[0_usize];
        send_node_removed(&mut sys, removed);
        procs.remove(0);
        step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    }

    check(&mut sys, &procs, &kv, false, false)
}

pub fn test_scale_up_down(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let mut rand = Pcg64::seed_from_u64(config.seed);
    let mut procs = sys.process_names();

    // insert random key-value pairs
    let keys_count = 1000;
    let mut kv = IndexMap::new();
    for _ in 0..keys_count {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(&sys, &mut rand);
        check_put(&mut sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }

    // add new N nodes to the system
    for i in 0..config.proc_count {
        let added = format!("{}", config.proc_count + i);
        add_node(&added, &mut sys, config);
        send_node_added(&mut sys, &added);
        procs.push(added);
        step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    }

    check(&mut sys, &procs, &kv, false, false)?;

    // remove new N nodes
    for _ in 0..config.proc_count {
        let removed = &procs[config.proc_count as usize];
        send_node_removed(&mut sys, removed);
        procs.remove(config.proc_count as usize);
        step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    }

    check(&mut sys, &procs, &kv, false, false)
}

const DISTRIBUTION_KEYS: usize = 50_000;

fn insert_distribution_keys(
    sys: &mut System,
    seed: u64,
) -> Result<IndexMap<String, String>, String> {
    let mut rand = Pcg64::seed_from_u64(seed);
    let mut kv = IndexMap::new();
    while kv.len() < DISTRIBUTION_KEYS {
        let k = random_string(8, &mut rand).to_uppercase();
        let v = random_string(8, &mut rand);
        let proc = random_proc(sys, &mut rand);
        check_put(sys, &proc, &k, &v, 100)?;
        kv.insert(k, v);
    }
    Ok(kv)
}

pub fn test_distribution(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let kv = insert_distribution_keys(&mut sys, config.seed)?;
    let procs = sys.process_names();
    check(&mut sys, &procs, &kv, false, true)
}

pub fn test_distribution_node_added(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    let kv = insert_distribution_keys(&mut sys, config.seed)?;
    let dist_before = key_distribution(&mut sys)?;

    // add new node to the system
    let added = sys.process_names().len().to_string();
    add_node(&added, &mut sys, config);
    send_node_added(&mut sys, &added);

    // run the system until key the distribution is stabilized
    let procs = sys.process_names();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    let dist_after = key_distribution(&mut sys)?;

    let target_moved_keys = (kv.len() as f64 / procs.len() as f64).round() as u64;
    check_moved_keys(&mut sys, &dist_before, &dist_after, target_moved_keys)?;

    check(&mut sys, &procs, &kv, false, true)
}

pub fn test_distribution_node_removed(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config, false);
    // Keep the topology choice independent of the number of generated keys.
    let removed = random_proc(&sys, &mut Pcg64::seed_from_u64(config.seed));
    let kv = insert_distribution_keys(&mut sys, config.seed)?;
    let dist_before = key_distribution(&mut sys)?;

    // remove a node from the system
    println!("Removed node: {removed}");
    let count = count_records(&mut sys, &removed)?;
    assume!(count > 0, "Node stores no records, bad distribution")?;
    send_node_removed(&mut sys, &removed);

    // run the system until key the distribution is stabilized
    let procs: Vec<String> = sys
        .process_names()
        .into_iter()
        .filter(|x| *x != removed)
        .collect();
    step_until_stabilized(&mut sys, &procs, kv.len() as u64, 100, 1000)?;
    let dist_after = key_distribution(&mut sys)?;

    let target_moved_keys = (kv.len() as f64 / (procs.len() + 1) as f64).round() as u64;
    check_moved_keys(&mut sys, &dist_before, &dist_after, target_moved_keys)?;

    check(&mut sys, &procs, &kv, false, true)
}

// UTILS ---------------------------------------------------------------------------------------------------------------

fn check_get(
    sys: &mut System,
    proc: &str,
    key: &str,
    expected: Option<&str>,
    max_steps: u32,
) -> TestResult {
    sys.send_local_message(proc, Message::json("GET", &GetReqMessage { key }));
    let res = sys.step_until_local_message_max_steps(proc, max_steps);
    assume!(res.is_ok(), format!("GET_RESP is not returned by {}", proc))?;
    let msgs = res.unwrap();
    let msg = msgs.first().unwrap();
    assume_eq!(msg.tip, "GET_RESP")?;
    let data: GetRespMessage = serde_json::from_str(&msg.data).unwrap();
    assume_eq!(data.key, key)?;
    assume!(
        data.value == expected,
        format!(
            "GET({key}) via {proc}: expected {expected:?}, got {:?}",
            data.value
        )
    )?;
    Ok(true)
}

fn wait_get_responses(
    sys: &mut System,
    proc: &str,
    mut expected: IndexMap<String, Option<String>>,
    max_steps: u32,
) -> TestResult {
    while !expected.is_empty() {
        let res = sys.step_until_local_message_max_steps(proc, max_steps);
        assume!(res.is_ok(), format!("GET_RESP is not returned by {proc}"))?;
        for msg in res.unwrap() {
            assume_eq!(msg.tip, "GET_RESP")?;
            let data: GetRespMessage = serde_json::from_str(&msg.data)
                .map_err(|error| format!("Invalid GET_RESP via {proc}: {error}"))?;
            let value = expected.shift_remove(data.key).ok_or_else(|| {
                format!("Unexpected or duplicate GET_RESP({}) via {proc}", data.key)
            })?;
            assume!(
                data.value == value.as_deref(),
                format!(
                    "GET({}) via {proc}: expected {value:?}, got {:?}",
                    data.key, data.value
                )
            )?;
        }
    }
    Ok(true)
}

fn check_put(sys: &mut System, proc: &str, key: &str, value: &str, max_steps: u32) -> TestResult {
    sys.send_local_message(proc, Message::json("PUT", &PutReqMessage { key, value }));
    let res = sys.step_until_local_message_max_steps(proc, max_steps);
    assume!(res.is_ok(), format!("PUT_RESP is not returned by {}", proc))?;
    let msgs = res.unwrap();
    let msg = msgs.first().unwrap();
    assume_eq!(msg.tip, "PUT_RESP")?;
    let data: PutRespMessage = serde_json::from_str(&msg.data).unwrap();
    assume_eq!(data.key, key)?;
    assume_eq!(data.value, value)?;
    Ok(true)
}

fn wait_put_responses(
    sys: &mut System,
    proc: &str,
    mut expected: IndexMap<String, String>,
    max_steps: u32,
) -> TestResult {
    while !expected.is_empty() {
        let res = sys.step_until_local_message_max_steps(proc, max_steps);
        assume!(res.is_ok(), format!("PUT_RESP is not returned by {}", proc))?;
        let msgs = res.unwrap();
        for msg in msgs {
            assume_eq!(msg.tip, "PUT_RESP")?;
            let data: PutRespMessage = serde_json::from_str(&msg.data)
                .map_err(|error| format!("Invalid PUT_RESP via {proc}: {error}"))?;
            let value = expected.shift_remove(data.key).ok_or_else(|| {
                format!("Unexpected or duplicate PUT_RESP({}) via {proc}", data.key)
            })?;
            assume!(
                data.value == value,
                format!(
                    "PUT({}) via {proc}: expected {value:?}, got {:?}",
                    data.key, data.value
                )
            )?;
        }
    }
    Ok(true)
}

fn check_delete(
    sys: &mut System,
    proc: &str,
    key: &str,
    expected: Option<&str>,
    max_steps: u32,
) -> TestResult {
    sys.send_local_message(proc, Message::json("DELETE", &DeleteReqMessage { key }));
    let res = sys.step_until_local_message_max_steps(proc, max_steps);
    assume!(
        res.is_ok(),
        format!("DELETE_RESP is not returned by {}", proc)
    )?;
    let msgs = res.unwrap();
    let msg = msgs.first().unwrap();
    assume_eq!(msg.tip, "DELETE_RESP")?;
    let data: DeleteRespMessage = serde_json::from_str(&msg.data).unwrap();
    assume_eq!(data.key, key)?;
    assume_eq!(data.value, expected)?;
    Ok(true)
}

fn wait_delete_responses(
    sys: &mut System,
    proc: &str,
    mut expected: IndexMap<String, Option<String>>,
    max_steps: u32,
) -> TestResult {
    while !expected.is_empty() {
        let res = sys.step_until_local_message_max_steps(proc, max_steps);
        assume!(
            res.is_ok(),
            format!("DELETE_RESP is not returned by {proc}")
        )?;
        for msg in res.unwrap() {
            assume_eq!(msg.tip, "DELETE_RESP")?;
            let data: DeleteRespMessage = serde_json::from_str(&msg.data)
                .map_err(|error| format!("Invalid DELETE_RESP via {proc}: {error}"))?;
            let value = expected.shift_remove(data.key).ok_or_else(|| {
                format!(
                    "Unexpected or duplicate DELETE_RESP({}) via {proc}",
                    data.key
                )
            })?;
            assume!(
                data.value == value.as_deref(),
                format!(
                    "DELETE({}) via {proc}: expected {value:?}, got {:?}",
                    data.key, data.value
                )
            )?;
        }
    }
    Ok(true)
}

fn send_node_added(sys: &mut System, added: &str) {
    for proc in sys.process_names() {
        if !sys.proc_node_is_crashed(&proc) {
            sys.send_local_message(
                &proc,
                Message::json("NODE_ADDED", &NodeMessage { id: added }),
            );
        }
    }
}

fn send_node_removed(sys: &mut System, removed: &str) {
    for proc in sys.process_names() {
        if !sys.proc_node_is_crashed(&proc) {
            sys.send_local_message(
                &proc,
                Message::json("NODE_REMOVED", &NodeMessage { id: removed }),
            );
        }
    }
}

fn count_records(sys: &mut System, proc: &str) -> Result<u64, String> {
    sys.send_local_message(proc, Message::json("COUNT_RECORDS", &EmptyMessage {}));
    let res = sys.step_until_local_message_max_steps(proc, 100);
    assume!(
        res.is_ok(),
        format!("COUNT_RECORDS_RESP is not returned by {}", proc)
    )?;
    let msgs = res.unwrap();
    let msg = msgs.first().unwrap();
    assume_eq!(msg.tip, "COUNT_RECORDS_RESP")?;
    let data: CountRecordsRespMessage = serde_json::from_str(&msg.data).unwrap();
    Ok(data.count)
}

fn dump_keys(sys: &mut System, proc: &str) -> Result<HashSet<String>, String> {
    sys.send_local_message(proc, Message::json("DUMP_KEYS", &EmptyMessage {}));
    let res = sys.step_until_local_message_max_steps(proc, 100);
    assume!(
        res.is_ok(),
        format!("DUMP_KEYS_RESP is not returned by {}", proc)
    )?;
    let msgs = res.unwrap();
    let msg = msgs.first().unwrap();
    assume_eq!(msg.tip, "DUMP_KEYS_RESP")?;
    let data: DumpKeysRespMessage = serde_json::from_str(&msg.data).unwrap();
    Ok(data.keys)
}

fn key_distribution(sys: &mut System) -> Result<HashMap<String, HashSet<String>>, String> {
    let mut dist = HashMap::new();
    for proc in sys.process_names() {
        dist.insert(proc.clone(), dump_keys(sys, &proc)?);
    }
    Ok(dist)
}

fn step_until_stabilized(
    sys: &mut System,
    procs: &[String],
    expected_keys: u64,
    steps_per_iter: u64,
    max_steps: u64,
) -> TestResult {
    let mut stabilized = false;
    let mut steps: u64 = 0;
    let mut counts = HashMap::new();
    let mut total_count: u64 = 0;
    for proc in procs.iter() {
        let count = count_records(sys, proc)?;
        counts.insert(proc, count);
        total_count += count;
    }

    while !stabilized && steps <= max_steps {
        sys.steps(steps_per_iter);
        steps += steps_per_iter;
        total_count = 0;
        let mut count_changed = false;
        for proc in procs.iter() {
            let count = count_records(sys, proc)?;
            if *counts.get(proc).unwrap() != count {
                count_changed = true;
            }
            counts.insert(proc, count);
            total_count += count;
        }
        if total_count == expected_keys && !count_changed {
            stabilized = true;
        }
    }

    assume!(
        stabilized,
        format!(
            "Keys distribution is not stabilized (keys observed = {}, expected = {})",
            total_count, expected_keys
        )
    )
}

fn check(
    sys: &mut System,
    procs: &[String],
    expected: &IndexMap<String, String>,
    check_values: bool,
    check_distribution: bool,
) -> TestResult {
    let mut stored_keys = HashSet::new();
    let mut proc_key_counts = Vec::new();
    for proc in procs.iter() {
        let proc_count = count_records(sys, proc)?;
        let proc_keys = dump_keys(sys, proc)?;
        assume_eq!(proc_keys.len() as u64, proc_count)?;
        stored_keys.extend(proc_keys);
        proc_key_counts.push(proc_count);
    }

    // all keys are stored
    assume!(
        expected.len() == stored_keys.len() && expected.keys().all(|k| stored_keys.contains(k)),
        "Stored keys do not mach expected"
    )?;

    // each key is stored on a single node
    assume!(
        proc_key_counts.iter().sum::<u64>() == stored_keys.len() as u64,
        "Keys are not stored on a single node"
    )?;

    // check values
    if check_values {
        println!("\nChecking values:");
        for proc in procs.iter() {
            for (k, v) in expected.iter() {
                check_get(sys, proc, k, Some(v), 100)?;
            }
        }
        println!("OK")
    }

    // check keys distribution
    if check_distribution {
        let target_count = (expected.len() as f64 / proc_key_counts.len() as f64).round();
        let max_count = *proc_key_counts.iter().max().unwrap();
        let min_count = *proc_key_counts.iter().min().unwrap();
        let deviations: Vec<f64> = proc_key_counts
            .iter()
            .map(|x| (target_count - *x as f64).abs() / target_count)
            .collect();
        let avg_deviation = deviations.iter().sum::<f64>() / proc_key_counts.len() as f64;
        let max_deviation = deviations
            .iter()
            .map(|x| R64::from_inner(*x))
            .max()
            .unwrap();
        println!("\nStored keys per node:");
        println!("  - total: {}", stored_keys.len());
        println!("  - target: {target_count}");
        println!("  - min: {min_count}");
        println!("  - max: {max_count}");
        println!("  - average deviation from target: {avg_deviation:.3}");
        println!("  - max deviation from target: {max_deviation:.3}");
        assume!(
            max_deviation <= 0.1,
            "Max deviation from target is above 10%"
        )?;
    }

    Ok(true)
}

fn check_moved_keys(
    sys: &mut System,
    before: &HashMap<String, HashSet<String>>,
    after: &HashMap<String, HashSet<String>>,
    target: u64,
) -> TestResult {
    let mut total_count_before = 0;
    let mut total_count_after = 0;
    let mut not_moved_count = 0;
    let empty = HashSet::new();
    for proc in sys.process_names() {
        let b = before.get(&proc).unwrap_or(&empty);
        total_count_before += b.len() as u64;
        let a = after.get(&proc).unwrap_or(&empty);
        total_count_after += a.len() as u64;
        let not_moved: HashSet<String> = a.intersection(b).cloned().collect();
        not_moved_count += not_moved.len() as u64;
    }
    assume_eq!(
        total_count_after,
        total_count_before,
        format!(
            "Number of stored keys has changed (before: {}, after: {})",
            total_count_before, total_count_after
        ),
    )?;
    let moved_count = total_count_before - not_moved_count;
    let deviation = (moved_count as f64 - target as f64) / target as f64;
    println!("\nMoved keys:");
    println!("  - target: {target}");
    println!("  - observed: {moved_count}");
    println!("  - deviation: {deviation:.3}");
    assume!(
        deviation <= 0.1,
        format!("Deviation from target is above 10%")
    )
}

#[cfg(test)]
mod memory_tests {
    use super::check_memory_growth;

    #[test]
    fn accepts_the_limit_and_rejects_one_more_byte_in_the_cluster() {
        let small = 765_301;
        let at_limit = 4 * small + 64 * 20_000;
        assert!(check_memory_growth(small, at_limit).is_ok());
        assert!(check_memory_growth(small, at_limit + 1).is_err());
        assert!(check_memory_growth(small, small).is_ok());
    }
}