//! Run each correctness test in a killable process. Student callbacks can block
//! inside Python, so an in-simulator deadline alone cannot enforce wall time.
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::path::PathBuf;
use std::process::Command;
use std::thread;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anysystem::python::PyProcessFactory;
use anysystem::test::TestResult;
use env_logger::Builder;
use indexmap::IndexMap;
use log::LevelFilter;

use crate::common::TestConfig;
use crate::tests::*;
use crate::tests_mc::*;
use crate::Args;

struct Test {
    name: &'static str,
    run: fn(&TestConfig) -> TestResult,
}

const TESTS: &[Test] = &[
    Test {
        name: "NORMAL",
        run: test_normal,
    },
    Test {
        name: "SENDER CRASH",
        run: test_sender_crash,
    },
    Test {
        name: "SENDER CRASH 2",
        run: test_sender_crash2,
    },
    Test {
        name: "TWO CRASHES",
        run: test_two_crashes,
    },
    Test {
        name: "TWO CRASHES 2",
        run: test_two_crashes2,
    },
    Test {
        name: "TWO CRASHES AFTER DELIVERY",
        run: test_two_crashes_after_delivery,
    },
    Test {
        name: "CAUSAL ORDER",
        run: test_causal_order,
    },
    Test {
        name: "QUORUM",
        run: test_quorum,
    },
    Test {
        name: "CHAOS MONKEY",
        run: test_chaos_monkey,
    },
    Test {
        name: "SCALABILITY",
        run: test_scalability,
    },
    Test {
        name: "MODEL CHECKING NORMAL DELIVERY",
        run: test_mc_normal_delivery,
    },
    Test {
        name: "MODEL CHECKING SENDER CRASH",
        run: test_mc_sender_crash,
    },
    Test {
        name: "MODEL CHECKING CAUSAL ORDER",
        run: test_mc_causal_order,
    },
];

pub fn worker(args: &Args) {
    let test = TESTS
        .iter()
        .find(|test| Some(test.name) == args.test_worker.as_deref())
        .unwrap();
    if args.debug {
        Builder::new()
            .filter(Some("anysystem"), LevelFilter::Debug)
            .format(|buf, record| writeln!(buf, "{}", record.args()))
            .init();
    }
    let factory = PyProcessFactory::new(&args.solution_path, "BroadcastProcess");
    let config = TestConfig {
        proc_factory: &factory,
        proc_count: if test.name.starts_with("MODEL CHECKING") {
            3
        } else {
            args.proc_count
        },
        seed: args.seed,
        monkeys: args.monkeys,
        debug: args.debug,
    };
    let result = (test.run)(&config);
    // Keep the structured result separate from student prints and debug traces.
    fs::write(
        args.test_result.as_ref().unwrap(),
        serde_json::to_vec(&result).unwrap(),
    )
    .unwrap();
}

struct ResultFile(PathBuf);

impl ResultFile {
    fn new() -> std::io::Result<Self> {
        let nonce = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!(
            "broadcast-test-{}-{nonce}.json",
            std::process::id()
        ));
        OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(&path)?;
        Ok(Self(path))
    }
}

impl Drop for ResultFile {
    fn drop(&mut self) {
        let _ = fs::remove_file(&self.0);
    }
}

fn run_one(args: &Args, test: &Test, suite_deadline: Instant) -> TestResult {
    let seconds = if test.name == "CHAOS MONKEY" {
        args.chaos_timeout
    } else if test.name == "SCALABILITY" {
        args.scalability_timeout
    } else if test.name.starts_with("MODEL CHECKING") {
        // The process deadline also catches a callback hanging inside the MC.
        args.mc_timeout
            .min(if test.name == "MODEL CHECKING NORMAL DELIVERY" {
                110
            } else if test.name == "MODEL CHECKING CAUSAL ORDER" {
                60
            } else {
                210
            })
    } else {
        args.test_timeout
    };
    let test_deadline = Instant::now() + Duration::from_secs(seconds);
    let deadline = test_deadline.min(suite_deadline);
    let timeout = if suite_deadline <= test_deadline {
        format!("suite time limit of {}s exceeded", args.suite_timeout)
    } else {
        format!("test time limit of {seconds}s exceeded")
    };
    if Instant::now() >= deadline {
        return Err(timeout);
    }
    let result_file = ResultFile::new().map_err(|e| format!("test worker error: {e}"))?;
    let mut command = Command::new(std::env::current_exe().map_err(|e| e.to_string())?);
    command
        .args([
            "--test-worker",
            test.name,
            "-i",
            &args.solution_path,
            "-s",
            &args.seed.to_string(),
            "-p",
            &args.proc_count.to_string(),
            "-m",
            &args.monkeys.to_string(),
            "--test-result",
        ])
        .arg(&result_file.0)
        .env("PYTHONUNBUFFERED", "1")
        // Python must see the seed before interpreter initialization.
        .env("PYTHONHASHSEED", args.seed.to_string());
    if args.debug {
        command.arg("-d");
    }
    let mut child = command
        .spawn()
        .map_err(|e| format!("test worker error: {e}"))?;
    loop {
        match child.try_wait() {
            Ok(Some(status)) => {
                if !status.success() {
                    return Err(format!("test worker exited with {status}"));
                }
                return serde_json::from_slice(
                    &fs::read(&result_file.0).map_err(|e| format!("test worker result: {e}"))?,
                )
                .map_err(|e| format!("test worker result: {e}"))?;
            }
            Ok(None) => {}
            Err(error) => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(format!("test worker error: {error}"));
            }
        }
        if Instant::now() >= deadline {
            let _ = child.kill();
            let _ = child.wait();
            return Err(timeout);
        }
        thread::sleep(Duration::from_millis(10));
    }
}

fn must_stop(result: &TestResult, keep_going: bool) -> bool {
    let Err(error) = result else {
        return false;
    };
    // Resource exhaustion and process failures always stop, including in local
    // diagnostic mode. A pure causal violation must not hide a later base error.
    if error.contains("time limit")
        || error.starts_with("causal model-checking ")
        || error.starts_with("test worker")
    {
        return true;
    }
    if keep_going {
        return false;
    }
    !error.strip_prefix("Violated ").is_some_and(|properties| {
        properties
            .split(", ")
            .all(|property| property == "CAUSAL ORDER")
    })
}

pub(crate) fn causal_mc_incomplete(name: &str, error: &str) -> bool {
    name == "MODEL CHECKING CAUSAL ORDER"
        && (error.contains("time limit") || error.starts_with("causal model-checking "))
}

pub fn run(args: &Args) -> IndexMap<String, TestResult> {
    let tests: Vec<_> = TESTS
        .iter()
        .filter(|test| {
            if let Some(name) = &args.test {
                test.name == name
            } else {
                !args.disable_mc_tests || !test.name.starts_with("MODEL CHECKING")
            }
        })
        .collect();
    let deadline = Instant::now() + Duration::from_secs(args.suite_timeout);
    let mut results = IndexMap::new();
    let mut stopped = None;
    let mut passed = 0;
    let mut causal_violated = false;
    for test in &tests {
        if let Some(failed_test) = stopped {
            println!("SKIPPED: {} (stopped after {failed_test})", test.name);
            continue;
        }
        if test.name == "MODEL CHECKING CAUSAL ORDER" && causal_violated && !args.keep_going {
            println!("SKIPPED: {} (CAUSAL ORDER already violated)", test.name);
            continue;
        }
        println!("\n--- {} ---\n", test.name);
        std::io::stdout().flush().unwrap();
        let result = run_one(args, test, deadline);
        match &result {
            Ok(_) => {
                println!("\nPASSED\n");
                passed += 1;
            }
            Err(error) => {
                causal_violated |= error.strip_prefix("Violated ").is_some_and(|properties| {
                    properties
                        .split(", ")
                        .any(|property| property == "CAUSAL ORDER")
                });
                println!("\nFAILED: {error}");
                if causal_mc_incomplete(test.name, error) {
                    println!("INCOMPLETE: causal model checking did not finish within its resource budget");
                }
                println!();
            }
        }
        if must_stop(&result, args.keep_going) {
            stopped = Some(test.name);
        }
        results.insert(test.name.to_string(), result);
    }
    if let Some(name) = &args.test {
        if tests.is_empty() {
            println!("FAILED: unknown test {name}");
        }
    } else {
        println!("-------------------------------------------------------------------------------");
        println!("\nPassed {passed} from {} tests\n", tests.len());
        if results.len() != tests.len() {
            println!("Skipped {} tests\n", tests.len() - results.len());
        }
        if results.values().any(Result::is_err) {
            println!("Failed tests:");
            for (name, result) in &results {
                if let Err(error) = result {
                    println!("- {name}: {error}");
                }
            }
            println!();
        }
    }
    results
}

#[cfg(test)]
mod tests {
    use super::must_stop;

    #[test]
    fn mc_timeout_always_stops_but_causal_violation_does_not() {
        assert!(must_stop(&Err("time limit of 100s exceeded".into()), false));
        assert!(must_stop(&Err("time limit of 100s exceeded".into()), true));
        assert!(!must_stop(&Err("Violated CAUSAL ORDER".into()), false));
        assert!(must_stop(
            &Err("Violated VALIDITY, CAUSAL ORDER".into()),
            false
        ));
        assert!(!must_stop(&Err("Violated VALIDITY".into()), true));
    }
}