#!/usr/bin/env python3
"""Talks to the server as an editor would, over stdin and stdout, and checks
what it answers. One server process for every case, so that a second check
in the same process is tested too.

    test/smoke.py        # builds nothing; run `dotnet build -c Release` first
"""

import json
import os
import subprocess
import sys
import tempfile

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SERVER = os.path.join(ROOT, "bin", "Release", "net10.0", "BjolangLsp.dll")


def frame(message):
    body = json.dumps(message).encode()
    return b"Content-Length: %d\r\n\r\n" % len(body) + body


def read(stream):
    length = None
    while True:
        line = stream.readline()
        if not line:
            raise EOFError("server closed stdout")
        line = line.strip()
        if not line:
            break
        name, _, value = line.decode().partition(":")
        if name.lower() == "content-length":
            length = int(value)
    return json.loads(stream.read(length))


class Server:
    def __init__(self):
        self.proc = subprocess.Popen(
            ["dotnet", SERVER], stdin=subprocess.PIPE, stdout=subprocess.PIPE
        )
        self.next_id = 0
        # Notifications read while waiting for something else: a check that
        # fell due can publish between a request and its answer.
        self.pending = []

    def send(self, method, params=None, request=False):
        message = {"jsonrpc": "2.0", "method": method, "params": params or {}}
        if request:
            self.next_id += 1
            message["id"] = self.next_id
        self.proc.stdin.write(frame(message))
        self.proc.stdin.flush()
        if not request:
            return None
        while True:
            answer = read(self.proc.stdout)
            if answer.get("id") == self.next_id:
                return answer
            self.pending.append(answer)

    def change(self, uri, text, version):
        self.send(
            "textDocument/didChange",
            {"textDocument": {"uri": uri, "version": version}, "contentChanges": [{"text": text}]},
        )

    def published(self, uri):
        """The diagnostics for `uri` among the notifications already read, as
        a list of their diagnostics lists, and takes them out."""
        mine = [
            m for m in self.pending
            if m.get("method") == "textDocument/publishDiagnostics" and m["params"]["uri"] == uri
        ]
        self.pending = [m for m in self.pending if m not in mine]
        return [m["params"]["diagnostics"] for m in mine]

    def next_diagnostics(self, uri):
        """The next diagnostics published for `uri`, waiting for them if
        none has come yet."""
        earlier = self.published(uri)
        if earlier:
            return earlier[0]
        while True:
            message = read(self.proc.stdout)
            if message.get("method") == "textDocument/publishDiagnostics" and message["params"]["uri"] == uri:
                return message["params"]["diagnostics"]
            self.pending.append(message)

    def diagnostics_after(self, method, uri, text=None):
        doc = {"uri": uri}
        if text is not None:
            doc.update(languageId="bjolang", version=1, text=text)
        self.published(uri)  # anything older is not the answer to this
        self.send(method, {"textDocument": doc})
        return self.next_diagnostics(uri)


failures = 0


def expect(name, ok, detail=""):
    global failures
    print(("ok: " if ok else "FAIL: ") + name + ("" if ok else f"\n  {detail}"))
    if not ok:
        failures += 1


BROKEN = """(import (std prelude))

(: main (-> (Vec string) int))
(defun (main args)
  (+ 1 "two"))
"""

FIXED = """(import (std prelude))

(: main (-> (Vec string) int))
(defun (main args)
  (+ 1 2))
"""

UNCLOSED = """(import (std prelude))

(defun (main args)
  (+ 1 2)
"""

SYMBOLS = """(import (std prelude))

(type (: Shape (Union (Circle int) (Square int) Empty)))
(type (: Point (Record (: x int) (: y int))))

(def/trait (Area %a)
  (: area (-> %a int)))

(impl (Area Shape)
  (defun (area s)
    (match s
      ((Circle r) (* 3 (* r r)))
      ((Square a) (* a a))
      (Empty 0))))

(def unit-size 1)

(: double (-> int int))
(defun (double n) (* 2 n))

(: main (-> (Vec string) int))
(defun (main args)
  (double (area (Circle unit-size)))
  (list-map double (list 1 2 3))
  0)
"""


def position_of(text, needle, nth=0, offset=1):
    """The LSP position `offset` characters into the `nth` occurrence of
    `needle` in `text`."""
    at = -1
    for _ in range(nth + 1):
        at = text.index(needle, at + 1)
    at += offset
    line = text.count("\n", 0, at)
    return {"line": line, "character": at - (text.rfind("\n", 0, at) + 1)}


def line_of(text, needle):
    return text.count("\n", 0, text.index(needle))


def names(symbols):
    return [(s["name"], names(s.get("children", []))) for s in symbols]


def symbols_and_definitions(server, tmp):
    path = os.path.join(tmp, "symbols.bjo")
    uri = "file://" + path
    with open(path, "w") as f:
        f.write(SYMBOLS)
    server.diagnostics_after("textDocument/didOpen", uri, SYMBOLS)

    found = server.send("textDocument/documentSymbol", {"textDocument": {"uri": uri}}, request=True)["result"]
    expect(
        "the outline",
        names(found)
        == [
            ("Shape", [("Circle", []), ("Square", []), ("Empty", [])]),
            ("Point", [("x", []), ("y", [])]),
            ("Area", [("area", [])]),
            ("impl (Area Shape)", [("area", [])]),
            ("unit-size", []),
            ("double", []),
            ("main", []),
        ],
        names(found),
    )

    def definition(text, needle, offset=1):
        answer = server.send(
            "textDocument/definition",
            {"textDocument": {"uri": uri}, "position": position_of(text, needle, 0, offset)},
            request=True,
        )
        return answer["result"]

    d = definition(SYMBOLS, "(double (area")
    expect(
        "a name defined in the file",
        d and d["uri"] == uri and d["range"]["start"]["line"] == line_of(SYMBOLS, "(defun (double"),
        d,
    )
    d = definition(SYMBOLS, "(Circle unit-size")
    expect(
        "a constructor",
        d and d["range"]["start"]["line"] == line_of(SYMBOLS, "(type (: Shape"),
        d,
    )
    d = definition(SYMBOLS, "(list-map")
    expect("a name from the prelude", d and d["uri"].endswith("/lib/std/prelude.bjo"), d)
    # `n` is a parameter: it goes to where the parameter is written.
    d = definition(SYMBOLS, "(* 2 n)", 5)
    expect(
        "a parameter goes to the parameter",
        d and d["range"]["start"] == position_of(SYMBOLS, "(defun (double n)", 0, 15),
        d,
    )

    def hover(text, needle, offset=1):
        answer = server.send(
            "textDocument/hover",
            {"textDocument": {"uri": uri}, "position": position_of(text, needle, 0, offset)},
            request=True,
        )
        result = answer["result"]
        return result["contents"]["value"] if result else None

    h = hover(SYMBOLS, "(double (area")
    expect(
        "hover: the declaration and the head",
        h and "(: double (-> int int))" in h and "(defun (double n))" in h,
        h,
    )
    h = hover(SYMBOLS, "(Circle unit-size")
    expect("hover: a constructor's case and type", h and "(Circle int)" in h and "A case of `Shape`" in h, h)
    h = hover(SYMBOLS, "(list-map")
    expect(
        "hover: the prelude's doc",
        h and "Applies a function to each element of a list." in h and "**Parameters**" in h,
        h,
    )
    h = hover(SYMBOLS, "(* 2 n)", 5)
    expect("hover: a parameter says whose", h and "*A parameter of `double`.*" in h, h)

    # Unsaved: the buffer gains a definition the file on disk does not have.
    changed = SYMBOLS + "\n(defun (triple n) (* 3 n))\n(def nine (triple 3))\n"
    server.send(
        "textDocument/didChange",
        {"textDocument": {"uri": uri, "version": 2}, "contentChanges": [{"text": changed}]},
    )
    d = definition(changed, "(triple 3)")
    expect(
        "the buffer, not the file",
        d and d["range"]["start"]["line"] == line_of(changed, "(defun (triple"),
        d,
    )
    found = server.send("textDocument/documentSymbol", {"textDocument": {"uri": uri}}, request=True)["result"]
    expect("the outline follows the buffer", [n for n, _ in names(found)][-2:] == ["triple", "nine"], names(found))

    # Half typed: the outline stays as it was rather than emptying.
    server.send(
        "textDocument/didChange",
        {"textDocument": {"uri": uri, "version": 3}, "contentChanges": [{"text": changed + "(defun (half"}]},
    )
    found = server.send("textDocument/documentSymbol", {"textDocument": {"uri": uri}}, request=True)["result"]
    expect(
        "a half-typed form is read as if closed",
        [n for n, _ in names(found)][-3:] == ["triple", "nine", "half"],
        names(found),
    )

    # A stray closer cannot be mended by closing: the last outline stays.
    server.send(
        "textDocument/didChange",
        {"textDocument": {"uri": uri, "version": 4}, "contentChanges": [{"text": changed + ")"}]},
    )
    found = server.send("textDocument/documentSymbol", {"textDocument": {"uri": uri}}, request=True)["result"]
    expect(
        "text that does not read keeps the last outline",
        [n for n, _ in names(found)][-3:] == ["triple", "nine", "half"],
        names(found),
    )

    server.diagnostics_after("textDocument/didClose", uri)


DOCUMENTED = """(import (std prelude))

(:doc #:file "documented.bjodoc")

(: squeeze (-> int int int int))
(defun (squeeze lo hi x) (max lo (min hi x)))

(: main (-> (Vec string) int))
(defun (main args) (squeeze 0 10 12))
"""

DOCUMENTED_DOC = '''(:doc squeeze
  (summary "@code{x}, kept @strong{between} @em{lo} and hi.")
  (arg lo "The least.")
  (arg hi "The most.")
  (arg x "The value.")
  (returns "One of: @itemlist{@item{@code{lo}} @item{@code{hi}} @item{x}}")
  (example """
    (squeeze 0 10 12)   ; 10
    """))
'''


def documented(server, tmp):
    path = os.path.join(tmp, "documented.bjo")
    uri = "file://" + path
    with open(path, "w") as f:
        f.write(DOCUMENTED)
    with open(os.path.join(tmp, "documented.bjodoc"), "w") as f:
        f.write(DOCUMENTED_DOC)
    found = server.diagnostics_after("textDocument/didOpen", uri, DOCUMENTED)
    expect("the documented file checks", found == [], found)

    answer = server.send(
        "textDocument/hover",
        {"textDocument": {"uri": uri}, "position": position_of(DOCUMENTED, "(squeeze 0")},
        request=True,
    )
    h = answer["result"]["contents"]["value"] if answer["result"] else None
    expect(
        "hover: a doc from a .bjodoc, its markup as Markdown",
        h
        and "`x`, kept **between** *lo* and hi." in h
        and "- `lo`: The least." in h
        and "**Returns** One of:\n- `lo`\n- `hi`\n- x" in h
        and "(squeeze 0 10 12)   ; 10" in h,
        h,
    )
    server.diagnostics_after("textDocument/didClose", uri)


COMPLETING = """(import (std prelude))

(defun (list-unwind xs) xs)

(: main (-> (Vec string) int))
(defun (main args)
  ;; list-u
  (list-u
"""


def completing(server, tmp):
    path = os.path.join(tmp, "completing.bjo")
    uri = "file://" + path
    with open(path, "w") as f:
        f.write(COMPLETING)
    server.diagnostics_after("textDocument/didOpen", uri, COMPLETING)

    def complete(needle, nth=0):
        at = position_of(COMPLETING, needle, nth, len(needle))
        answer = server.send(
            "textDocument/completion", {"textDocument": {"uri": uri}, "position": at}, request=True
        )
        return at, answer["result"]["items"]

    at, items = complete("(list-u", 0)
    labels = {i["label"]: i for i in items}
    expect(
        "completion: the file's own name and the prelude's",
        "list-unwind" in labels and "list-unfold" in labels,
        sorted(labels)[:20],
    )
    expect(
        "completion: only names starting with the word's first letter, in either case",
        all(l[0] in "lL" for l in labels),
        sorted(labels)[:20],
    )
    own = labels.get("list-unwind", {})
    expect(
        "completion: the item replaces the whole word",
        own.get("textEdit", {}).get("range", {}).get("start") == {"line": at["line"], "character": at["character"] - 6},
        own,
    )
    expect(
        "completion: the file's own names sort first",
        own.get("sortText", "").startswith("1") and labels["list-unfold"]["sortText"].startswith("2"),
        (own.get("sortText"), labels.get("list-unfold", {}).get("sortText")),
    )

    resolved = server.send("completionItem/resolve", labels["list-unfold"], request=True)["result"]
    doc = resolved.get("documentation", {}).get("value", "")
    expect("completion: resolve adds the doc", "A list generated from a state." in doc, doc)

    _, items = complete(";; list-u")
    expect("completion: none in a comment", items == [], [i["label"] for i in items][:10])

    # A builtin: no source to define it, but a type and a doc.
    seq_text = COMPLETING.replace("(list-u\n", "(seq-unfold\n")
    server.send(
        "textDocument/didChange",
        {"textDocument": {"uri": uri, "version": 2}, "contentChanges": [{"text": seq_text}]},
    )
    answer = server.send(
        "textDocument/hover",
        {"textDocument": {"uri": uri}, "position": position_of(seq_text, "(seq-unfold")},
        request=True,
    )
    h = answer["result"]["contents"]["value"] if answer["result"] else None
    expect(
        "hover: a builtin's type and doc",
        h and "(: seq-unfold (->" in h and "A sequence generated from a state, lazily." in h,
        h,
    )
    server.diagnostics_after("textDocument/didClose", uri)


LOCALS = """(import (std prelude))

(: total (-> (List int) int int))
(defun (total xs (: bonus int))
  (def base 10)
  (let ((doubled (list-map (fun (x) (* 2 x)) xs))
        (count 0))
    (match doubled
      ((Cons first rest) (+ first bonus base count))
      (Nil (loop (:for y xs)
                 (:acc sum (summing (+ y bonus)))
                 => sum)))))

(: shadow (-> int int))
(defun (shadow v)
  (let ((v (+ v 1)))
    v))

(: typing (-> int int))
(defun (typing amount)
  (let ((amplified (* amount 2)))
    (+ am
"""


def locals_(server, tmp):
    path = os.path.join(tmp, "locals.bjo")
    uri = "file://" + path
    with open(path, "w") as f:
        f.write(LOCALS)
    server.diagnostics_after("textDocument/didOpen", uri, LOCALS)

    def ask(method, needle, offset=1, nth=0):
        at = position_of(LOCALS, needle, nth, offset)
        return server.send(method, {"textDocument": {"uri": uri}, "position": at}, request=True)["result"]

    def goes_to(needle, target, target_offset, offset=1, nth=0):
        d = ask("textDocument/definition", needle, offset, nth)
        return d and d["range"]["start"] == position_of(LOCALS, target, 0, target_offset), d

    cases = [
        ("a parameter", "bonus base", "(: bonus int)", 3),
        ("a def in the body", "base count", "(def base", 5),
        ("a let binding", "doubled\n", "(doubled (list-map", 1),
        ("a fun's parameter", "x)) xs", "(fun (x)", 6),
        ("a match pattern's variable", "first bonus", "(Cons first", 6),
        ("a loop variable, from its :acc", "y bonus", "(:for y", 6),
        ("an accumulator, from =>", "sum)))))", "(:acc sum", 6),
        ("the inner of two of a name", "v))\n", "((v (+ v 1)))", 2),
        ("the outer, in the inner's value", "v 1)", "(shadow v)", 8),
    ]
    for name, needle, target, target_offset in cases:
        ok, d = goes_to(needle, target, target_offset, 0)
        expect("local definition: " + name, ok, d)

    h = ask("textDocument/hover", "bonus base", 0)
    h = h and h["contents"]["value"]
    expect("local hover: a typed parameter shows its type", h and "(: bonus int)" in h and "A parameter of `total`" in h, h)
    h = ask("textDocument/hover", "first bonus", 0)
    h = h and h["contents"]["value"]
    expect("local hover: a pattern shows the pattern", h and "(Cons first rest)" in h and "`match` pattern" in h, h)

    # Completion while the form is being typed: the locals of `typing`, and
    # nothing of `total`'s.
    at = position_of(LOCALS, "(+ am", 0, 5)
    items = server.send("textDocument/completion", {"textDocument": {"uri": uri}, "position": at}, request=True)["result"]["items"]
    labels = {i["label"]: i for i in items}
    expect(
        "local completion: what is in scope, first",
        labels.get("amount", {}).get("sortText", "").startswith("0")
        and labels.get("amplified", {}).get("sortText", "").startswith("0"),
        sorted((l, i["sortText"]) for l, i in labels.items())[:10],
    )
    expect("local completion: another function's locals are not", "doubled" not in labels and "bonus" not in labels, sorted(labels)[:10])
    doc = labels.get("amplified", {}).get("documentation", {}).get("value", "")
    expect("local completion: its doc comes along", "(amplified (* amount 2))" in doc, doc)
    server.diagnostics_after("textDocument/didClose", uri)


CHAIN_MACROS = """;; Two macros, one expanding to the other.
(import (std syntax-match))

(def/macro (twice form inject compare)
  (syntax-match form
    ((_ e) #'(both ,e ,e))))

(def/macro (both form inject compare)
  (syntax-match form
    ((_ a b) #'(let ((tmp ,a)) (+ tmp ,b)))))
"""

CHAIN = """(import "chain-macros.bjo")

(: main (-> (Vec string) int))
(defun (main args)
  (twice 21))
"""


def expanding(server, tmp):
    with open(os.path.join(tmp, "chain-macros.bjo"), "w") as f:
        f.write(CHAIN_MACROS)
    path = os.path.join(tmp, "chain.bjo")
    uri = "file://" + path
    with open(path, "w") as f:
        f.write(CHAIN)
    found = server.diagnostics_after("textDocument/didOpen", uri, CHAIN)
    expect("expansion: a file with an imported macro module checks", found == [], found)
    expect(
        "expansion: the import was built by the installed compiler",
        os.path.exists(os.path.join(tmp, "chain-macros.dll")),
        os.listdir(tmp),
    )

    def expand(needle, **flags):
        at = position_of(CHAIN, needle, 0, 2)
        answer = server.send(
            "bjolang/expandMacro", dict(textDocument={"uri": uri}, position=at, **flags), request=True
        )
        return answer["result"]

    r = expand("(twice", full=False)
    expect("expansion: one step", r and r["macro"] == "twice" and r["text"] == "(both 21 21)", r)
    r = expand("(twice")
    expect("expansion: all the way", r and r["text"] == "(let ((tmp 21)) (+ tmp 21))", r)
    r = expand("(twice", hygienic=True)
    expect("expansion: hygienic names kept", r and "tmp__" in r["text"], r)
    r = expand("(defun (main", full=False)
    expect("expansion: none where there is no macro call", r is None, r)

    # An unsaved change is expanded as it stands, without waiting for the pause.
    changed = CHAIN.replace("(twice 21)", "(twice (twice 1))")
    server.change(uri, changed, 2)
    at = position_of(changed, "(twice (twice", 0, 2)
    r = server.send("bjolang/expandMacro", {"textDocument": {"uri": uri}, "position": at}, request=True)["result"]
    expect(
        "expansion: of the buffer, nested calls expanded too",
        r and r["text"].count("(+ tmp") == 3 and "twice" not in r["text"] and "both" not in r["text"],
        r,
    )
    server.diagnostics_after("textDocument/didClose", uri)


with tempfile.TemporaryDirectory() as tmp:
    path = os.path.join(tmp, "smoke.bjo")
    uri = "file://" + path
    server = Server()

    answer = server.send("initialize", {"processId": None, "rootUri": None, "capabilities": {}}, request=True)
    expect("initialize answers capabilities", "capabilities" in answer.get("result", {}), answer)
    server.send("initialized")

    with open(path, "w") as f:
        f.write(BROKEN)
    found = server.diagnostics_after("textDocument/didOpen", uri, BROKEN)
    expect("a type error is reported", len(found) >= 1 and found[0]["severity"] == 1, found)
    expect(
        "on the line it is on",
        any(d["range"]["start"]["line"] == 4 for d in found),
        json.dumps(found, indent=1),
    )
    for d in found:
        print("   ", d["range"]["start"], d["message"].splitlines()[0])

    # Saved: the editor sends the change, writes the file, and says so.
    server.change(uri, FIXED, 2)
    with open(path, "w") as f:
        f.write(FIXED)
    found = server.diagnostics_after("textDocument/didSave", uri)
    expect("the fixed file has none", found == [], json.dumps(found, indent=1))

    # Not saved: the check follows the buffer once typing pauses, while the
    # file on disk stays fixed.
    server.change(uri, UNCLOSED, 3)
    found = server.next_diagnostics(uri)
    expect("a parse error in an unsaved buffer is reported", len(found) >= 1, found)
    for d in found:
        print("   ", d["range"]["start"], d["message"].splitlines()[0])

    # A burst of typing is checked once, after it, not once per change.
    for version in range(4, 9):
        server.change(uri, FIXED if version % 2 else BROKEN, version)
    found = server.next_diagnostics(uri)
    expect("a burst is checked as its last text", len(found) >= 1, found)
    import time
    time.sleep(1.0)
    server.send("textDocument/documentSymbol", {"textDocument": {"uri": uri}}, request=True)
    expect("and only once", server.published(uri) == [], "more than one check published")

    found = server.diagnostics_after("textDocument/didClose", uri)
    expect("closing clears them", found == [], found)

    symbols_and_definitions(server, tmp)
    documented(server, tmp)
    completing(server, tmp)
    locals_(server, tmp)
    expanding(server, tmp)

    answer = server.send("shutdown", request=True)
    expect("shutdown answers", "result" in answer, answer)
    server.send("exit")
    expect("exit ends the process with 0", server.proc.wait(timeout=10) == 0)

sys.exit(1 if failures else 0)