use std::cell::Cell;
use std::collections::HashSet;
use std::rc::Rc;
use std::time::Duration;

use serde_json::Value;
use sugars::boxed;

use anysystem::logger::LogEntry;
use anysystem::test::TestResult;
use anysystem::Message;

use anysystem::mc::predicates::{collects, goals, invariants, prunes};
use anysystem::mc::strategies::Bfs;
use anysystem::mc::strategy::{GoalFn, InvariantFn, PruneFn, StrategyConfig};
use anysystem::mc::{McState, ModelChecker};

use crate::common::{build_system, BroadcastMessage, TestConfig};

// The causal test explores two broadcasts on three processes. The state budget
// is checked as an invariant, so exhausting it cannot silently pass a pruned search.
const CAUSAL_MC_STATE_LIMIT: usize = 100_000;

pub fn test_mc_normal_delivery(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config);
    let proc_names = sys.process_names();
    let text = "0:Hello";
    sys.send_local_message(
        proc_names[0].as_str(),
        Message::json("SEND", &BroadcastMessage { text }),
    );
    let goal = goals::all_goals(
        proc_names
            .iter()
            .map(|name| goals::got_n_local_messages(name, name, 1))
            .collect::<Vec<GoalFn>>(),
    );
    let strategy_config = StrategyConfig::default()
        .goal(goal)
        .prune(prunes::any_prune(vec![
            prunes::state_depth(10),
            mc_prune_proc_permutations(&proc_names[1..]),
            // Prune states with more than 2 messages received from any process
            mc_prune_msg_per_proc_limit(&proc_names, 2),
        ]))
        .invariant(invariants::all_invariants(vec![
            mc_invariant(proc_names.clone(), text.to_string()),
            invariants::time_limit(Duration::from_secs(100)),
        ]));
    let mut mc = ModelChecker::new(&sys);
    let res = mc.run::<Bfs>(strategy_config);
    if let Err(err) = res {
        err.print_trace();
        Err(err.message())
    } else {
        Ok(true)
    }
}

pub fn test_mc_sender_crash(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config);
    let proc_names = sys.process_names();
    let text = "0:Hello";
    sys.send_local_message(
        proc_names[0].as_str(),
        Message::json("SEND", &BroadcastMessage { text }),
    );
    let goal = goals::all_goals(
        proc_names
            .iter()
            .map(|name| goals::got_n_local_messages(name, name, 1))
            .collect::<Vec<GoalFn>>(),
    );

    let strategy_config = StrategyConfig::default()
        .prune(prunes::any_prune(vec![
            prunes::state_depth(4),
            mc_prune_proc_permutations(&proc_names[1..]),
        ]))
        .goal(goal)
        .invariant(invariants::all_invariants(vec![
            mc_invariant(proc_names.clone(), text.to_string()),
            invariants::time_limit(Duration::from_secs(100)),
        ]))
        .collect(collects::any_collect(
            proc_names[1..]
                .iter()
                .map(|proc| collects::got_n_local_messages(proc, proc, 1))
                .collect(),
        ));

    let mut mc = ModelChecker::new(&sys);
    let res = mc.run::<Bfs>(strategy_config);
    let intermediate_states = res
        .map_err(|err| {
            err.print_trace();
            err.message()
        })?
        .collected_states;
    if intermediate_states.is_empty() {
        return Err("no states collected after first stage".to_string());
    }

    // Crash first node in the list
    let left_proc_names = proc_names[1..].to_vec();
    let goal = goals::all_goals(
        left_proc_names
            .iter()
            .map(|name| goals::got_n_local_messages(name, name, 1))
            .collect::<Vec<GoalFn>>(),
    );
    let strategy_config = StrategyConfig::default()
        .goal(goal)
        .invariant(invariants::all_invariants(vec![
            mc_invariant(left_proc_names, text.to_string()),
            invariants::time_limit(Duration::from_secs(100)),
        ]))
        .prune(prunes::any_prune(vec![
            prunes::state_depth(6),
            mc_prune_proc_permutations(&proc_names[1..]),
            // Prune states with more than 4 messages received from any process
            mc_prune_msg_per_proc_limit(&proc_names, 4),
        ]));
    let res = mc.run_from_states_with_change::<Bfs>(strategy_config, intermediate_states, |sys| {
        sys.crash_node(proc_names[0].clone());
    });
    if let Err(err) = res {
        err.print_trace();
        Err(err.message())
    } else {
        Ok(true)
    }
}

// Send "a" from P2, then send "b" from P1 after P1 delivers "a",
// which gives the causal order "a" -> "b", then check that all processes
// deliver both messages and deliver "a" before "b".
pub fn test_mc_causal_order(config: &TestConfig) -> TestResult {
    let mut sys = build_system(config);
    sys.send_local_message("2", Message::json("SEND", &BroadcastMessage { text: "a" }));
    let proc_names = sys.process_names();
    let report_error = |err: anysystem::mc::McError| {
        err.print_trace();
        err.message()
    };
    let checked = Rc::new(Cell::new(0));
    let make_strategy = |texts: Vec<&'static str>| {
        let goal_procs = proc_names.clone();
        let invariant_procs = proc_names.clone();
        let checked = checked.clone();
        StrategyConfig::default()
            // Stop this branch when everyone has delivered the messages
            // sent in this stage. If events end before that, the test fails.
            .goal(boxed!(move |state: &McState| {
                goal_procs
                    .iter()
                    .all(|proc| {
                        texts
                            .iter()
                            .all(|text| mc_delivery_index(state, proc, text).is_some())
                    })
                    .then(|| "all messages have been delivered by every process".to_string())
            }))
            // Since P1 sends "b" after it delivers "a",
            // every process must also deliver "a" before "b".
            .invariant(invariants::all_invariants(vec![
                boxed!(move |state: &McState| {
                    for proc in &invariant_procs {
                        if let Some(b) = mc_delivery_index(state, proc, "b") {
                            if !mc_delivery_index(state, proc, "a").is_some_and(|a| a < b) {
                                return Err("Violated CAUSAL ORDER".to_string());
                            }
                        }
                    }
                    Ok(())
                }),
                boxed!(move |_: &McState| {
                    let next = checked.get() + 1;
                    checked.set(next);
                    if next > CAUSAL_MC_STATE_LIMIT {
                        Err(format!(
                            "causal model-checking state limit of {CAUSAL_MC_STATE_LIMIT} exceeded"
                        ))
                    } else {
                        Ok(())
                    }
                }),
                invariants::time_limit(Duration::from_secs(55)),
            ]))
    };
    let mut mc = ModelChecker::new(&sys);

    // Explore all schedules of "a" before P1 sends "b". The first stage
    // checks delivery of "a" on every branch and retains every state where
    // P1 can send "b" after delivering "a".
    let states = mc
        .run::<Bfs>(make_strategy(vec!["a"]).collect(boxed!(|state: &McState| {
            mc_delivery_index(state, "1", "a").is_some()
        })))
        .map_err(&report_error)?
        .collected_states;
    if states.is_empty() {
        return Err("No state with delivery of a on process 1 was reached".to_string());
    }

    // Explore every continuation after sending "b" from each collected state.
    // Resource exhaustion is reported as incomplete rather than pruning paths.
    mc.run_from_states_with_change::<Bfs>(make_strategy(vec!["a", "b"]), states, |sys| {
        sys.send_local_message(
            "1",
            "1",
            Message::json("SEND", &BroadcastMessage { text: "b" }),
        );
    })
    .map(|_| true)
    .map_err(report_error)
}

fn mc_delivery_index(state: &McState, proc: &str, text: &str) -> Option<usize> {
    state.node_states[proc].proc_states[proc]
        .local_outbox
        .iter()
        .position(|msg| {
            msg.tip == "DELIVER"
                && serde_json::from_str::<Value>(&msg.data).unwrap()["text"] == text
        })
}

fn mc_prune_proc_permutations(equivalent_procs: &[String]) -> PruneFn {
    let equivalent_procs = equivalent_procs.to_vec();
    boxed!(move |state| {
        let proc_names = HashSet::<String>::from_iter(equivalent_procs.clone().into_iter());
        let mut used_proc_names = HashSet::<String>::new();
        let mut waiting_for_proc = 0;
        for entry in &state.trace {
            match entry {
                LogEntry::McMessageReceived { src: proc, .. }
                | LogEntry::McTimerFired { proc, .. } => {
                    if used_proc_names.contains(proc) || !proc_names.contains(proc) {
                        continue;
                    }
                    if equivalent_procs[waiting_for_proc] != *proc {
                        return Some(
                            "state is the same as another state with renumerated processes"
                                .to_owned(),
                        );
                    }
                    used_proc_names.insert(proc.clone());
                    waiting_for_proc += 1;
                }
                _ => {}
            }
        }
        None
    })
}

fn mc_prune_msg_per_proc_limit(proc_names: &[String], limit: usize) -> PruneFn {
    prunes::events_limit_per_proc(
        |entry: &LogEntry, proc: &String| match entry {
            LogEntry::McMessageReceived { src, .. } => src == proc,
            _ => false,
        },
        proc_names.to_owned(),
        limit,
    )
}

fn mc_invariant(proc_names: Vec<String>, sent_message: String) -> InvariantFn {
    boxed!(move |state: &McState| {
        for name in &proc_names {
            let outbox = &state.node_states[name].proc_states[name].local_outbox;
            let mut message_data = HashSet::new();
            for message in outbox {
                let data: Value = serde_json::from_str(&message.data).unwrap();
                let message = data["text"].as_str().unwrap().to_string();

                if message_data.contains(&message) {
                    return Err("Violated NO DUPLICATION".to_owned());
                }
                message_data.insert(message.clone());
                if message != sent_message {
                    return Err("Violated NO CREATION".to_owned());
                }
            }
        }
        Ok(())
    })
}