mod common;
mod optimization;
mod runner;
mod shared_state;
mod tests;
mod tests_mc;

use anysystem::test::TestResult;
use clap::Parser;
use indexmap::IndexMap;
use std::collections::HashSet;
use std::path::PathBuf;

/// Broadcast Homework Tests
#[derive(Parser, Debug)]
#[clap(about, long_about = None)]
struct Args {
    /// Path to Python file with solution
    #[clap(long = "impl", short = 'i', default_value = "solution/broadcast.py")]
    solution_path: String,

    /// Test to run (optional)
    #[clap(long = "test", short)]
    test: Option<String>,

    /// Print execution trace
    #[clap(long, short)]
    debug: bool,

    /// Random seed used in tests
    #[clap(long, short, default_value = "2023")]
    seed: u64,

    /// Number of processes
    #[clap(long, short, default_value = "5")]
    proc_count: u64,

    /// Number of chaos monkey runs
    #[clap(long, short, default_value = "10")]
    monkeys: u32,

    /// Do not run model checking tests
    #[clap(long)]
    disable_mc_tests: bool,

    /// Optional scalable solution (defaults to broadcast_scalable.py beside --impl)
    #[clap(long)]
    optimized_impl: Option<String>,

    /// Repetitions per OPTIMIZATION scenario and system size
    #[clap(long, default_value = "10", value_parser = clap::value_parser!(u32).range(1..=100))]
    optimization_runs: u32,

    /// Shared wall-clock budget in seconds for both implementations in OPTIMIZATION
    #[clap(long, default_value = "90", value_parser = clap::value_parser!(u64).range(1..=90))]
    optimization_timeout: u64,

    #[clap(long, hide = true)]
    optimization_worker: bool,

    /// Wall-clock limit in seconds per ordinary test
    #[clap(long, default_value = "15", value_parser = clap::value_parser!(u64).range(1..=300))]
    test_timeout: u64,

    /// Wall-clock limit in seconds for SCALABILITY (up to 50 processes)
    #[clap(long, default_value = "60", value_parser = clap::value_parser!(u64).range(1..=300))]
    scalability_timeout: u64,

    /// Wall-clock limit in seconds for all CHAOS MONKEY runs
    #[clap(long, default_value = "30", value_parser = clap::value_parser!(u64).range(1..=300))]
    chaos_timeout: u64,

    /// External MC test limit (NORMAL DELIVERY is capped at 110s, CAUSAL ORDER at 60s)
    #[clap(long, default_value = "210", value_parser = clap::value_parser!(u64).range(1..=210))]
    mc_timeout: u64,

    /// Total wall-clock budget in seconds for correctness tests
    #[clap(long, default_value = "300", value_parser = clap::value_parser!(u64).range(1..=300))]
    suite_timeout: u64,

    /// Continue after failures for local diagnostics; resource limits and crashes still stop
    #[clap(long)]
    keep_going: bool,

    #[clap(long, hide = true, requires = "test-result")]
    test_worker: Option<String>,

    #[clap(long, hide = true)]
    test_result: Option<PathBuf>,
}

fn main() {
    let args = Args::parse();
    if args.test_worker.is_some() {
        runner::worker(&args);
        return;
    }
    if args.optimization_worker {
        optimization::worker(&args);
        return;
    }
    if args.test.as_deref() == Some("OPTIMIZATION") {
        optimization::run(&args);
        return;
    }
    if !shared_state::validate_or_report(&args.solution_path, &["BroadcastProcess"]) {
        return;
    }
    let results = runner::run(&args);
    if args.test.is_none() {
        let score = score(results);
        optimization::run_if_present(&args, score);
        println!("SCORE: {score}\n");
    }
}

fn score(results: IndexMap<String, TestResult>) -> f32 {
    let mut violated = HashSet::new();
    let mut causal_incomplete = false;
    for (name, result) in results {
        if let Err(e) = result {
            if e.starts_with("Violated") {
                for prop in e.replace("Violated ", "").split(", ") {
                    violated.insert(prop.to_string());
                }
            } else if runner::causal_mc_incomplete(&name, &e) {
                // Keep partial credit separate from a discovered violation.
                causal_incomplete = true;
            } else {
                return 0.;
            }
        }
    }
    let violated: HashSet<&str> = violated.iter().map(|s| s.as_str()).collect();
    let base_props = HashSet::from([
        "NO DUPLICATION",
        "NO CREATION",
        "VALIDITY",
        "UNIFORM AGREEMENT",
    ]);
    if violated.is_disjoint(&base_props) {
        if !violated.is_empty() {
            5.
        } else if causal_incomplete {
            6.
        } else {
            7.
        }
    } else {
        0.
    }
}

#[cfg(test)]
mod score_tests {
    use super::score;
    use indexmap::IndexMap;

    #[test]
    fn incomplete_causal_mc_gets_partial_credit() {
        for error in [
            "time limit of 55s exceeded",
            "test time limit of 60s exceeded",
            "suite time limit of 300s exceeded",
            "causal model-checking state limit of 100000 exceeded",
        ] {
            let mut results = IndexMap::new();
            results.insert("MODEL CHECKING CAUSAL ORDER".to_string(), Err(error.into()));
            assert_eq!(score(results), 6., "{error}");
        }
    }

    #[test]
    fn discovered_violation_takes_precedence_over_incomplete_search() {
        for (error, expected) in [
            ("Violated CAUSAL ORDER", 5.),
            ("Violated VALIDITY", 0.),
            ("Violated VALIDITY, CAUSAL ORDER", 0.),
        ] {
            for incomplete_first in [false, true] {
                let mut entries = vec![
                    ("CAUSAL ORDER", error),
                    ("MODEL CHECKING CAUSAL ORDER", "time limit of 55s exceeded"),
                ];
                if incomplete_first {
                    entries.reverse();
                }
                let results = entries
                    .into_iter()
                    .map(|(name, error)| (name.to_string(), Err(error.to_string())))
                    .collect();
                assert_eq!(score(results), expected, "{error}");
            }
        }
    }

    #[test]
    fn other_failures_do_not_get_partial_credit() {
        for (name, error) in [
            ("NORMAL", "test time limit of 15s exceeded"),
            (
                "MODEL CHECKING NORMAL DELIVERY",
                "time limit of 100s exceeded",
            ),
            (
                "MODEL CHECKING CAUSAL ORDER",
                "test worker exited with exit code: 1",
            ),
            (
                "MODEL CHECKING CAUSAL ORDER",
                "nothing left to do to reach the goal",
            ),
        ] {
            let mut results = IndexMap::new();
            results.insert(name.to_string(), Err(error.to_string()));
            assert_eq!(score(results), 0., "{name}: {error}");
        }
    }
}