//! `wai-quantum lsp` driven as an editor drives it: the binary started, a
//! scripted session of protocol messages written to its standard input, and
//! every reply read back from its standard output and checked.

use serde_json::{Value, json};
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::sync::mpsc::{Receiver, channel};
use std::time::Duration;

struct Editor {
    child: Child,
    stdin: ChildStdin,
    rx: Receiver<Value>,
    next: i64,
    notes: Vec<Value>,
    home: std::path::PathBuf,
}

fn read_message(r: &mut impl BufRead) -> Option<Value> {
    let mut len = 0usize;
    loop {
        let mut line = String::new();
        if r.read_line(&mut line).ok()? == 0 {
            return None;
        }
        let line = line.trim_end();
        if line.is_empty() {
            break;
        }
        if let Some(v) = line.strip_prefix("Content-Length:") {
            len = v.trim().parse().ok()?;
        }
    }
    let mut buf = vec![0u8; len];
    r.read_exact(&mut buf).ok()?;
    serde_json::from_slice(&buf).ok()
}

impl Editor {
    fn start(name: &str) -> Editor {
        let home = std::env::temp_dir().join(format!("wai-quantum-lsp-{name}-{}", std::process::id()));
        std::fs::create_dir_all(&home).unwrap();
        let mut child = Command::new(env!("CARGO_BIN_EXE_wai-quantum"))
            .args(["lsp", "--demo-key", "--memory-gb", "1"])
            .env("XDG_CONFIG_HOME", &home)
            .env_remove("WAI_QUANTUM_KEY")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .expect("the binary starts");
        let stdin = child.stdin.take().unwrap();
        let mut stdout = BufReader::new(child.stdout.take().unwrap());
        let (tx, rx) = channel();
        std::thread::spawn(move || {
            while let Some(m) = read_message(&mut stdout) {
                if tx.send(m).is_err() {
                    break;
                }
            }
        });
        Editor { child, stdin, rx, next: 1, notes: Vec::new(), home }
    }

    fn write(&mut self, msg: &Value) {
        let body = msg.to_string();
        write!(self.stdin, "Content-Length: {}\r\n\r\n{body}", body.len()).unwrap();
        self.stdin.flush().unwrap();
    }

    fn recv(&mut self) -> Value {
        self.rx.recv_timeout(Duration::from_secs(120)).expect("the server answers within two minutes")
    }

    /// A request's whole reply; notifications that come first are kept.
    fn request(&mut self, method: &str, params: Value) -> Value {
        let id = self.next;
        self.next += 1;
        self.write(&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }));
        loop {
            let m = self.recv();
            if m.get("id") == Some(&json!(id)) && (m.get("result").is_some() || m.get("error").is_some()) {
                return m;
            }
            self.notes.push(m);
        }
    }

    fn result(&mut self, method: &str, params: Value) -> Value {
        let r = self.request(method, params);
        assert!(r.get("error").is_none(), "{method}: {r}");
        r["result"].clone()
    }

    fn notify(&mut self, method: &str, params: Value) {
        self.write(&json!({ "jsonrpc": "2.0", "method": method, "params": params }));
    }

    /// The next notification that `pick` takes.
    fn note(&mut self, pick: impl Fn(&Value) -> bool) -> Value {
        if let Some(i) = self.notes.iter().position(&pick) {
            return self.notes.remove(i);
        }
        loop {
            let m = self.recv();
            if pick(&m) {
                return m;
            }
            self.notes.push(m);
        }
    }

    fn diagnostics(&mut self, uri: &str, version: i64) -> Vec<Value> {
        let n = self.note(|m| m["method"] == "textDocument/publishDiagnostics" && m["params"]["uri"] == uri && m["params"]["version"] == version);
        n["params"]["diagnostics"].as_array().unwrap().clone()
    }

    fn exit(mut self) -> i32 {
        self.notify("exit", json!(null));
        let status = self.child.wait().unwrap();
        let _ = std::fs::remove_dir_all(&self.home);
        status.code().unwrap_or(-1)
    }
}

fn at(line: u64, character: u64) -> Value {
    json!({ "line": line, "character": character })
}

const BAD: &str = "OPENQASM 3.0;\nqubit[1] q;\nfoo q[0];\n";
const GOOD: &str = "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nqubit[2] q;\nbit[2] c;\ngate bell a, b {\n  h a;\n  cx a, b;\n}\nx q[0];\nx q[0];\nbell q[0], q[1];\nc = measure q;\n";

#[test]
fn an_editor_s_session_from_initialize_to_exit() {
    let mut e = Editor::start("session");
    let init = e.result(
        "initialize",
        json!({
            "processId": null,
            "rootUri": null,
            "capabilities": { "general": { "positionEncodings": ["utf-16"] }, "textDocument": { "completion": { "completionItem": { "snippetSupport": true } } } },
            "initializationOptions": { "wai-quantum": { "shots": 200 } },
        }),
    );
    let caps = &init["capabilities"];
    assert_eq!(caps["positionEncoding"], "utf-16");
    assert_eq!(caps["textDocumentSync"]["change"], 1);
    for c in ["hoverProvider", "definitionProvider", "referencesProvider", "documentSymbolProvider", "foldingRangeProvider", "documentFormattingProvider", "documentHighlightProvider"] {
        assert_eq!(caps[c], true, "{c}");
    }
    assert_eq!(caps["renameProvider"]["prepareProvider"], true);
    assert_eq!(caps["executeCommandProvider"]["commands"], json!(["wai-quantum.run", "wai-quantum.plan", "wai-quantum.critique", "wai-quantum.shortcuts", "wai-quantum.receipt"]));
    assert_eq!(init["serverInfo"]["name"], "wai-quantum");
    e.notify("initialized", json!({}));

    // The reader's refusal, with its code, on the very text it names.
    let uri = "file:///work/bell.qasm";
    e.notify("textDocument/didOpen", json!({ "textDocument": { "uri": uri, "languageId": "openqasm", "version": 1, "text": BAD } }));
    let d = e.diagnostics(uri, 1);
    assert_eq!(d.len(), 1, "{d:?}");
    assert_eq!((d[0]["severity"].clone(), d[0]["message"].clone(), d[0]["source"].clone(), d[0]["code"].clone()), (json!(1), json!("unknown gate foo"), json!("wai-quantum"), json!("Q103")));
    assert_eq!(d[0]["range"], json!({ "start": at(2, 0), "end": at(2, 3) }));

    // A good program with two gates that cancel: a warning on each line.
    e.notify("textDocument/didChange", json!({ "textDocument": { "uri": uri, "version": 2 }, "contentChanges": [{ "text": GOOD }] }));
    let d = e.diagnostics(uri, 2);
    let cancels: Vec<&Value> = d.iter().filter(|x| x["code"] == "cancels").collect();
    assert_eq!(cancels.len(), 2, "{d:?}");
    assert_eq!((cancels[0]["severity"].clone(), cancels[0]["range"]["start"]["line"].clone(), cancels[1]["range"]["start"]["line"].clone()), (json!(2), json!(8), json!(9)));

    // Completion: the qubit register after a gate's name; gates, with placeholders, at a statement's start.
    let items = e.result("textDocument/completion", json!({ "textDocument": { "uri": uri }, "position": at(10, 5) }));
    assert_eq!(items["items"][0]["label"], "q");
    let items = e.result("textDocument/completion", json!({ "textDocument": { "uri": uri }, "position": at(12, 0) }));
    let all = items["items"].as_array().unwrap();
    let rz = all.iter().find(|i| i["label"] == "rz").unwrap();
    assert_eq!((rz["insertText"].clone(), rz["insertTextFormat"].clone()), (json!("rz(${1:theta})"), json!(2)));
    assert!(all.iter().any(|i| i["label"] == "bell" && i["kind"] == 3));

    // Hover: a standard gate's exact matrix; a user gate's, and its cost; a register's size.
    let h = e.result("textDocument/hover", json!({ "textDocument": { "uri": uri }, "position": at(5, 2) }));
    let text = h["contents"]["value"].as_str().unwrap();
    assert!(text.contains("[ 1/√2  1/√2  ]") && text.contains("Exact"), "{text}");
    assert_eq!(h["range"], json!({ "start": at(5, 2), "end": at(5, 3) }));
    let b = e.result("textDocument/hover", json!({ "textDocument": { "uri": uri }, "position": at(10, 1) }));
    let text = b["contents"]["value"].as_str().unwrap();
    assert!(text.contains("a gate this program defines, on 2 qubits") && text.contains("computed in double precision") && text.contains("two-qubit gates 1"), "{text}");
    let q = e.result("textDocument/hover", json!({ "textDocument": { "uri": uri }, "position": at(2, 9) }));
    assert!(q["contents"]["value"].as_str().unwrap().contains("qubit register, 2 qubits"));

    // Definition, references, highlights.
    let def = e.result("textDocument/definition", json!({ "textDocument": { "uri": uri }, "position": at(10, 1) }));
    assert_eq!(def, json!({ "uri": uri, "range": { "start": at(4, 5), "end": at(4, 9) } }));
    let refs = e.result("textDocument/references", json!({ "textDocument": { "uri": uri }, "position": at(10, 1), "context": { "includeDeclaration": true } }));
    assert_eq!(refs.as_array().unwrap().len(), 2);
    let refs = e.result("textDocument/references", json!({ "textDocument": { "uri": uri }, "position": at(10, 1), "context": { "includeDeclaration": false } }));
    assert_eq!(refs, json!([{ "uri": uri, "range": { "start": at(10, 0), "end": at(10, 4) } }]));
    let hl = e.result("textDocument/documentHighlight", json!({ "textDocument": { "uri": uri }, "position": at(5, 4) }));
    assert_eq!(hl.as_array().unwrap().iter().map(|x| x["kind"].clone()).collect::<Vec<_>>(), [json!(3), json!(2), json!(2)]);

    // Rename: the declaration and every use; a standard gate refused, with why.
    let p = e.result("textDocument/prepareRename", json!({ "textDocument": { "uri": uri }, "position": at(10, 2) }));
    assert_eq!(p, json!({ "range": { "start": at(10, 0), "end": at(10, 4) }, "placeholder": "bell" }));
    let edit = e.result("textDocument/rename", json!({ "textDocument": { "uri": uri }, "position": at(10, 2), "newName": "pair" }));
    let edits = edit["changes"][uri].as_array().unwrap();
    assert_eq!(edits.len(), 2);
    assert!(edits.iter().all(|x| x["newText"] == "pair"));
    let no = e.request("textDocument/rename", json!({ "textDocument": { "uri": uri }, "position": at(10, 2), "newName": "cx" }));
    assert!(no["error"]["message"].as_str().unwrap().contains("standard gate"), "{no}");
    let no = e.request("textDocument/prepareRename", json!({ "textDocument": { "uri": uri }, "position": at(5, 2) }));
    assert_eq!(no["error"]["code"], -32803);

    // The outline, the folds and the lenses.
    let syms = e.result("textDocument/documentSymbol", json!({ "textDocument": { "uri": uri } }));
    assert_eq!(syms.as_array().unwrap().iter().map(|s| s["name"].clone()).collect::<Vec<_>>(), [json!("q"), json!("c"), json!("bell")]);
    assert_eq!(syms[2]["children"].as_array().unwrap().len(), 2);
    assert_eq!(syms[2]["selectionRange"], json!({ "start": at(4, 5), "end": at(4, 9) }));
    let folds = e.result("textDocument/foldingRange", json!({ "textDocument": { "uri": uri } }));
    assert!(folds.as_array().unwrap().contains(&json!({ "startLine": 4, "endLine": 6 })), "{folds}");
    let lenses = e.result("textDocument/codeLens", json!({ "textDocument": { "uri": uri } }));
    let lenses = lenses.as_array().unwrap();
    assert_eq!(lenses.iter().map(|l| l["command"]["title"].clone()).collect::<Vec<_>>(), [json!("Run"), json!("Plan"), json!("Critique"), json!("Shortcuts"), json!("Receipt")]);
    assert_eq!(lenses[0]["command"], json!({ "title": "Run", "command": "wai-quantum.run", "arguments": [uri] }));
    assert_eq!(lenses[0]["range"]["start"], at(0, 0));

    // The commands, run here through the engine; a run's receipt verifies.
    let run = e.result("workspace/executeCommand", json!({ "command": "wai-quantum.run", "arguments": [uri] }));
    assert_eq!((run["ok"].clone(), run["json"]["method"].clone(), run["json"]["request"]["shots"].clone()), (json!(true), json!("stabilizer"), json!(200)));
    assert!(run["json"]["result"]["counts"].as_object().unwrap().keys().all(|k| k == "00" || k == "11"));
    assert!(wai_quantum::engine::job::verify(&run["json"]).unwrap().0);
    assert!(run["text"].as_str().unwrap().contains("receipt     verifies"), "{run}");
    let shown = e.note(|m| m["method"] == "window/showMessage");
    assert!(shown["params"]["message"].as_str().unwrap().starts_with("Run: stabilizer"), "{shown}");
    let plan = e.result("workspace/executeCommand", json!({ "command": "wai-quantum.plan", "arguments": [{ "uri": uri }] }));
    assert_eq!(plan["json"]["method"], "stabilizer");
    let crit = e.result("workspace/executeCommand", json!({ "command": "wai-quantum.critique", "arguments": [uri] }));
    assert!(crit["json"]["smells"].as_array().unwrap().iter().any(|s| s["kind"] == "cancels"));
    let short = e.result("workspace/executeCommand", json!({ "command": "wai-quantum.shortcuts", "arguments": [uri] }));
    assert_eq!(short["json"]["ok"], true);
    let receipt = e.result("workspace/executeCommand", json!({ "command": "wai-quantum.receipt", "arguments": [uri] }));
    let text = receipt["text"].as_str().unwrap();
    assert!(text.starts_with("receipt     verifies\nsigned by   did:key:"), "{text}");
    assert!(wai_quantum::engine::job::verify(&receipt["json"]).unwrap().0);
    let no = e.request("workspace/executeCommand", json!({ "command": "wai-quantum.nothing", "arguments": [uri] }));
    assert!(no["error"]["message"].as_str().unwrap().contains("the commands are"));

    // Formatting: one edit to the whole document, none for a formatted one, a refusal with its line.
    let messy = "file:///work/messy.qasm";
    e.notify("textDocument/didOpen", json!({ "textDocument": { "uri": messy, "languageId": "openqasm", "version": 1, "text": "OPENQASM 3.0;qubit[2]q;\nh q[0];cx q[0],q[1];\n" } }));
    let f = e.result("textDocument/formatting", json!({ "textDocument": { "uri": messy }, "options": { "tabSize": 4, "insertSpaces": true } }));
    assert_eq!(f, json!([{ "range": { "start": at(0, 0), "end": at(2, 0) }, "newText": "OPENQASM 3.0;\nqubit[2] q;\nh q[0];\ncx q[0], q[1];\n" }]));
    let f = e.result("textDocument/formatting", json!({ "textDocument": { "uri": uri }, "options": { "tabSize": 2, "insertSpaces": true } }));
    assert_eq!(f, json!([]));
    e.notify("textDocument/didChange", json!({ "textDocument": { "uri": messy, "version": 2 }, "contentChanges": [{ "text": "gate g a {\n  h a;\n" }] }));
    let no = e.request("textDocument/formatting", json!({ "textDocument": { "uri": messy }, "options": {} }));
    assert!(no["error"]["message"].as_str().unwrap().contains("line 1: `{` is not closed"), "{no}");

    // A change by range, as an editor that sends increments would.
    e.notify("textDocument/didChange", json!({ "textDocument": { "uri": messy, "version": 3 }, "contentChanges": [{ "range": { "start": at(1, 2), "end": at(1, 3) }, "text": "x" }] }));
    let _ = e.diagnostics(messy, 3);
    let hv = e.result("textDocument/hover", json!({ "textDocument": { "uri": messy }, "position": at(1, 2) }));
    assert!(hv["contents"]["value"].as_str().unwrap().contains("Pauli X"), "{hv}");

    // Positions count UTF-16 units: after `π`, column 10 is `q`.
    let pi = "file:///work/pi.qasm";
    e.notify("textDocument/didOpen", json!({ "textDocument": { "uri": pi, "languageId": "openqasm", "version": 1, "text": "OPENQASM 3.0;\nqubit[1] q;\nrz(π / 2) q[0];\n" } }));
    let hv = e.result("textDocument/hover", json!({ "textDocument": { "uri": pi }, "position": at(2, 10) }));
    assert!(hv["contents"]["value"].as_str().unwrap().contains("qubit register, 1 qubit"), "{hv}");
    assert_eq!(hv["range"], json!({ "start": at(2, 10), "end": at(2, 11) }));

    // A generator is read at the editor's inputs: a configuration change reads it again.
    let generator = "file:///work/generator.qasm";
    e.notify("textDocument/didOpen", json!({ "textDocument": { "uri": generator, "languageId": "openqasm", "version": 1, "text": "OPENQASM 3.0;\ninput uint n;\nqubit[n] q;\nh q[3];\n" } }));
    assert!(e.diagnostics(generator, 1).iter().all(|d| d["severity"] != 1));
    e.notify("workspace/didChangeConfiguration", json!({ "settings": { "wai-quantum": { "inputs": { "n": 3 } } } }));
    let d = e.diagnostics(generator, 1);
    assert_eq!(d[0]["message"], "q[3] is outside q, which has 3 (read at n = 3)");

    // Closing a document clears its diagnostics.
    e.notify("textDocument/didClose", json!({ "textDocument": { "uri": generator } }));
    let n = e.note(|m| m["method"] == "textDocument/publishDiagnostics" && m["params"]["uri"] == generator);
    assert_eq!(n["params"]["diagnostics"], json!([]));
    let no = e.request("textDocument/hover", json!({ "textDocument": { "uri": generator }, "position": at(0, 0) }));
    assert_eq!(no["error"]["code"], -32602);

    assert_eq!(e.result("shutdown", json!(null)), Value::Null);
    assert_eq!(e.exit(), 0);
}

#[test]
fn a_request_before_initialize_is_refused_and_an_exit_without_shutdown_fails() {
    let mut e = Editor::start("early");
    let r = e.request("textDocument/hover", json!({ "textDocument": { "uri": "file:///x.qasm" }, "position": at(0, 0) }));
    assert_eq!(r["error"]["code"], -32002);
    let r = e.request("initialize", json!({ "capabilities": { "general": { "positionEncodings": ["utf-8", "utf-16"] } } }));
    assert_eq!(r["result"]["capabilities"]["positionEncoding"], "utf-8");
    let r = e.request("textDocument/nothing", json!({ "textDocument": { "uri": "file:///x.qasm" } }));
    assert_eq!(r["error"]["code"], -32602);
    let r = e.request("workspace/nothing", json!({}));
    assert_eq!(r["error"]["code"], -32601);
    assert_eq!(e.exit(), 1);
}