//! `wai-quantum optimize --checkpoint`: a run stopped and started again ends
//! with the same bytes as one that never stopped, whether it stops by
//! `--slice` or is killed while it runs. Nothing here leaves this machine.

use std::path::{Path, PathBuf};
use std::process::{Command, Output, Stdio};
use std::time::{Duration, Instant};

fn bin() -> &'static str {
    env!("CARGO_BIN_EXE_wai-quantum")
}

/// A scratch directory of the test's own, and the program it optimises:
/// `n` qubits, a layer of `ry` on each, a chain of `cx`, and another layer.
fn scratch(tag: &str, n: usize) -> (PathBuf, PathBuf) {
    let dir = std::env::temp_dir().join(format!("wai-q-opt-{tag}-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let mut src = String::from("OPENQASM 3.0;\ninclude \"stdgates.inc\";\n");
    for k in 0..2 * n {
        src += &format!("input float t{k};\n");
    }
    src += &format!("qubit[{n}] q;\n");
    for k in 0..n {
        src += &format!("ry(t{k}) q[{k}];\n");
    }
    for k in 0..n - 1 {
        src += &format!("cx q[{k}], q[{}];\n", k + 1);
    }
    for k in 0..n {
        src += &format!("ry(t{}) q[{k}];\n", n + k);
    }
    let prog = dir.join("p.qasm");
    std::fs::write(&prog, src).unwrap();
    (dir, prog)
}

/// The command's arguments: every parameter from 0.1 up, the Hamiltonian
/// Z0 Z1 and X on the last qubit, Adam, `steps` steps, signed with the demo
/// key so that the same answer is the same bytes.
fn args(prog: &Path, n: usize, steps: u32) -> Vec<String> {
    let mut a = vec!["optimize".to_string(), prog.display().to_string()];
    for k in 0..2 * n {
        a.push("--param".into());
        a.push(format!("t{k}=0.{}", k % 9 + 1));
    }
    for t in ["Z0 Z1".to_string(), format!("0.5 X{}", n - 1)] {
        a.push("--term".into());
        a.push(t);
    }
    for x in ["--optimizer", "adam", "--steps", &steps.to_string(), "--json", "--demo-key"] {
        a.push(x.into());
    }
    a
}

fn run(a: &[String], extra: &[&str]) -> Output {
    Command::new(bin()).args(a).args(extra).env_remove("WAI_QUANTUM_KEY").output().unwrap()
}

#[test]
fn a_run_taken_in_slices_ends_with_the_bytes_of_one_that_never_stopped() {
    let (dir, prog) = scratch("slices", 3);
    let a = args(&prog, 3, 40);
    let whole = run(&a, &[]);
    assert!(whole.status.success(), "{}", String::from_utf8_lossy(&whole.stderr));
    let cp = dir.join("run.json");
    let cps = cp.display().to_string();
    let mut slices = 0;
    let last = loop {
        slices += 1;
        let o = run(&a, &["--checkpoint", &cps, "--slice", "7"]);
        match o.status.code() {
            Some(2) => {
                let err = String::from_utf8_lossy(&o.stderr);
                assert!(err.contains("the same command again goes on") && o.stdout.is_empty(), "{err}");
                assert!(cp.exists(), "a stopped run keeps its state");
            }
            Some(0) => break o,
            other => panic!("{other:?}: {}", String::from_utf8_lossy(&o.stderr)),
        }
        assert!(slices < 20, "the run never ended");
    };
    // 41 steps, seven a slice: five slices stop, the sixth ends.
    assert_eq!(slices, 6);
    assert_eq!(last.stdout, whole.stdout, "the same bytes as a run that never stopped");
    assert!(!cp.exists(), "the checkpoint is removed once the run ends");
    // A checkpoint of another request is refused, and left as it is.
    let o = run(&a, &["--checkpoint", &cps, "--slice", "3"]);
    assert_eq!(o.status.code(), Some(2));
    let other = args(&prog, 3, 41);
    let o = run(&other, &["--checkpoint", &cps]);
    assert_eq!(o.status.code(), Some(1));
    assert!(String::from_utf8_lossy(&o.stderr).contains("it is of another request"), "{}", String::from_utf8_lossy(&o.stderr));
    assert!(cp.exists());
    let _ = std::fs::remove_dir_all(&dir);
}

#[test]
fn a_run_killed_while_it_runs_goes_on_to_the_bytes_of_one_that_never_stopped() {
    // Big enough that a run outlasts its first checkpoint by seconds.
    let (n, steps) = (16, 600);
    let (dir, prog) = scratch("killed", n);
    let a = args(&prog, n, steps);
    let cp = dir.join("run.json");
    let cps = cp.display().to_string();
    let mut child = Command::new(bin()).args(&a).args(["--checkpoint", &cps]).env_remove("WAI_QUANTUM_KEY").stdout(Stdio::null()).stderr(Stdio::null()).spawn().unwrap();
    let t = Instant::now();
    // Killed once it has kept a state some steps in, as a crash or a power cut would stop it.
    let at = loop {
        if let Ok(text) = std::fs::read_to_string(&cp)
            && let Ok(v) = serde_json::from_str::<serde_json::Value>(&text)
            && v["next"].as_u64().is_some_and(|k| k >= 2)
        {
            break v["next"].as_u64().unwrap();
        }
        assert!(child.try_wait().unwrap().is_none(), "the run ended before it was killed: make it longer");
        assert!(t.elapsed() < Duration::from_secs(120), "no checkpoint was kept");
        std::thread::sleep(Duration::from_millis(20));
    };
    child.kill().unwrap();
    child.wait().unwrap();
    assert!(at < u64::from(steps), "killed partway");
    let resumed = run(&a, &["--checkpoint", &cps]);
    assert!(resumed.status.success(), "{}", String::from_utf8_lossy(&resumed.stderr));
    assert!(String::from_utf8_lossy(&resumed.stderr).contains("going on from step"), "{}", String::from_utf8_lossy(&resumed.stderr));
    let whole = run(&a, &[]);
    assert!(whole.status.success());
    assert_eq!(resumed.stdout, whole.stdout, "the same bytes as a run that never stopped");
    let _ = std::fs::remove_dir_all(&dir);
}