#!/usr/bin/env python3

import argparse
import json
import os
import select
import subprocess
import time
from pathlib import Path


def frame(message: dict) -> bytes:
    body = json.dumps(message, separators=(",", ":")).encode("utf-8")
    return f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body


def raw_frame(body: bytes) -> bytes:
    return f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body


def content_type_frame(message: dict) -> bytes:
    body = json.dumps(message, separators=(",", ":")).encode("utf-8")
    header = (
        f"Content-Length: {len(body)}\r\n"
        "Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n"
    ).encode("ascii")
    return header + body


def read_messages(data: bytes) -> list[dict]:
    messages = []
    offset = 0
    while offset < len(data):
        header_end = data.find(b"\r\n\r\n", offset)
        if header_end < 0:
            raise RuntimeError("Incomplete LSP response header")

        header = data[offset:header_end].decode("ascii")
        fields = dict(line.split(": ", 1) for line in header.split("\r\n"))
        length = int(fields["Content-Length"])
        body_start = header_end + 4
        body_end = body_start + length
        if body_end > len(data):
            raise RuntimeError("Incomplete LSP response body")

        messages.append(json.loads(data[body_start:body_end].decode("utf-8")))
        offset = body_end
    return messages


def response_by_id(data: bytes, request_id: int | str) -> dict:
    return next(message for message in read_messages(data)
                if message.get("id") == request_id)


def run_raw(server: Path, payload: bytes) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(
        [str(server), "--core-server-fixture"],
        input=payload,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=5,
        check=False,
    )


def run_batch(server: Path, messages: list[dict]) -> subprocess.CompletedProcess[bytes]:
    return run_raw(server, b"".join(frame(message) for message in messages))


class LspClient:
    def __init__(self, server: Path) -> None:
        self.proc = subprocess.Popen(
            [str(server), "--core-server-fixture"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        self.buffer = b""

    def send(self, message: dict) -> None:
        if self.proc.stdin is None:
            raise RuntimeError("Fixture stdin is unavailable")
        self.proc.stdin.write(frame(message))
        self.proc.stdin.flush()

    def read(self, timeout: float = 1.0) -> dict:
        if self.proc.stdout is None:
            raise RuntimeError("Fixture stdout is unavailable")

        deadline = time.monotonic() + timeout
        while True:
            header_end = self.buffer.find(b"\r\n\r\n")
            if header_end >= 0:
                header = self.buffer[:header_end].decode("ascii")
                fields = dict(line.split(": ", 1)
                              for line in header.split("\r\n"))
                length = int(fields["Content-Length"])
                body_start = header_end + 4
                body_end = body_start + length
                if len(self.buffer) >= body_end:
                    body = self.buffer[body_start:body_end]
                    self.buffer = self.buffer[body_end:]
                    return json.loads(body.decode("utf-8"))

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("Timed out waiting for LSP response")

            ready, _, _ = select.select(
                [self.proc.stdout.fileno()], [], [], remaining)
            if not ready:
                continue

            chunk = os.read(self.proc.stdout.fileno(), 4096)
            if not chunk:
                raise RuntimeError("Fixture closed stdout before a response")
            self.buffer += chunk

    def close_input(self) -> int:
        if self.proc.stdin is not None:
            self.proc.stdin.close()
        return self.proc.wait(timeout=5)

    def kill(self) -> None:
        if self.proc.poll() is None:
            self.proc.kill()
            self.proc.wait(timeout=5)


def assert_lifecycle(server: Path) -> None:
    before_init = run_batch(server, [
        {"jsonrpc": "2.0", "id": 1, "method": "shutdown"},
        {"jsonrpc": "2.0", "method": "exit"},
    ])
    if before_init.returncode != 1:
        raise RuntimeError("shutdown before initialize should exit with code 1")
    if response_by_id(before_init.stdout, 1)["error"]["code"] != -32002:
        raise RuntimeError("shutdown before initialize should return ServerNotInitialized")

    repeated = run_batch(server, [
        {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
        {"jsonrpc": "2.0", "id": 2, "method": "initialize", "params": {}},
        {"jsonrpc": "2.0", "id": 3, "method": "shutdown"},
        {"jsonrpc": "2.0", "method": "exit"},
    ])
    if repeated.returncode != 0:
        raise RuntimeError("shutdown followed by exit should return code 0")
    if response_by_id(repeated.stdout, 2)["error"]["code"] != -32600:
        raise RuntimeError("repeated initialize should return InvalidRequest")
    if response_by_id(repeated.stdout, 3).get("result", "missing") is not None:
        raise RuntimeError("shutdown should return a null result")

    direct_exit = run_batch(
        server, [{"jsonrpc": "2.0", "method": "exit"}])
    if direct_exit.returncode != 1:
        raise RuntimeError("exit before shutdown should return code 1")

    client = LspClient(server)
    try:
        client.send({
            "jsonrpc": "2.0",
            "id": 11,
            "method": "initialize",
            "params": {},
        })
        if client.read().get("id") != 11:
            raise RuntimeError("missing initialize response")

        client.send({"jsonrpc": "2.0", "id": 12, "method": "shutdown"})
        if client.read().get("result", "missing") is not None:
            raise RuntimeError("shutdown should return a null result")

        time.sleep(0.1)
        if client.proc.poll() is not None:
            raise RuntimeError("server exited before receiving exit notification")

        client.send({"jsonrpc": "2.0", "method": "exit"})
        if client.close_input() != 0:
            raise RuntimeError("orderly lifecycle should return code 0")
    finally:
        client.kill()


def assert_json_rpc_errors(server: Path) -> None:
    payload = b"".join([
        raw_frame(b'{"jsonrpc":"2.0",'),
        raw_frame(b"[]"),
        raw_frame(b'{"jsonrpc":"1.0","id":7,"method":"initialize"}'),
        frame({
            "jsonrpc": "2.0",
            "id": 8,
            "method": "initialize",
            "params": {},
        }),
        frame({"jsonrpc": "2.0", "id": 10, "method": "shutdown"}),
        frame({"jsonrpc": "2.0", "method": "exit"}),
    ])
    result = run_raw(server, payload)
    if result.returncode != 0:
        raise RuntimeError("JSON-RPC error sequence should shut down cleanly")

    responses = read_messages(result.stdout)
    if responses[0].get("id", "missing") is not None:
        raise RuntimeError(
            f"ParseError response should use a null id: {responses[0]!r}")
    if responses[0].get("error", {}).get("code") != -32700:
        raise RuntimeError("Malformed JSON should return ParseError")
    if responses[1].get("id", "missing") is not None:
        raise RuntimeError("InvalidRequest response should use a null id")
    if responses[1].get("error", {}).get("code") != -32600:
        raise RuntimeError("A non-object JSON value should return InvalidRequest")
    if response_by_id(result.stdout, 7).get("error", {}).get("code") != -32600:
        raise RuntimeError("jsonrpc other than 2.0 should return InvalidRequest")
    client = LspClient(server)
    try:
        client.send({
            "jsonrpc": "2.0",
            "id": 11,
            "method": "initialize",
            "params": {},
        })
        client.read()
        client.send({"jsonrpc": "2.0", "id": 9, "method": "missing/method"})
        if client.read().get("error", {}).get("code") != -32601:
            raise RuntimeError("Unknown request method should return MethodNotFound")
        client.send({"jsonrpc": "2.0", "id": 12, "method": "shutdown"})
        client.read()
        client.send({"jsonrpc": "2.0", "method": "exit"})
        if client.close_input() != 0:
            raise RuntimeError("unknown method sequence should shut down cleanly")
    finally:
        client.kill()


def assert_cancellation_and_failures(server: Path) -> None:
    client = LspClient(server)
    try:
        client.send({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {},
        })
        if client.read().get("id") != 1:
            raise RuntimeError("missing initialize response")

        client.send({"jsonrpc": "2.0", "id": 2, "method": "test/block"})
        client.send({"jsonrpc": "2.0", "id": "2", "method": "test/block"})
        client.send({
            "jsonrpc": "2.0",
            "method": "$/cancelRequest",
            "params": {"id": 2},
        })
        cancelled = client.read(timeout=1)
        if cancelled.get("id") != 2 or \
                cancelled.get("error", {}).get("code") != -32800:
            raise RuntimeError("integer request id should be cancelled independently")

        client.send({
            "jsonrpc": "2.0",
            "method": "$/cancelRequest",
            "params": {"id": "2"},
        })
        string_cancelled = client.read(timeout=1)
        if string_cancelled.get("id") != "2" or \
                string_cancelled.get("error", {}).get("code") != -32800:
            raise RuntimeError("string request id should be cancelled independently")

        client.send({"jsonrpc": "2.0", "id": 3, "method": "test/throw"})
        failed = client.read(timeout=1)
        if failed.get("id") != 3 or \
                failed.get("error", {}).get("code") != -32603:
            raise RuntimeError("request exceptions should return InternalError")

        client.send({"jsonrpc": "2.0", "id": 4, "method": "test/block"})
        client.send({"jsonrpc": "2.0", "id": 4, "method": "test/block"})
        duplicate = client.read(timeout=1)
        if duplicate.get("id") != 4 or \
                duplicate.get("error", {}).get("code") != -32600:
            raise RuntimeError("duplicate active request id should return InvalidRequest")

        client.send({
            "jsonrpc": "2.0",
            "method": "$/cancelRequest",
            "params": {"id": 4},
        })
        cancelled_original = client.read(timeout=1)
        if cancelled_original.get("id") != 4 or \
                cancelled_original.get("error", {}).get("code") != -32800:
            raise RuntimeError("duplicate id must not replace the original request")

        client.send({"jsonrpc": "2.0", "method": "test/throwNotification"})
        client.send({"jsonrpc": "2.0", "id": 5, "method": "shutdown"})
        shutdown = client.read(timeout=1)
        if shutdown.get("id") != 5 or shutdown.get("result", "missing") is not None:
            raise RuntimeError("notification exception should not stop the server")

        client.send({"jsonrpc": "2.0", "method": "exit"})
        if client.close_input() != 0:
            raise RuntimeError("fixture should exit cleanly after cancellation tests")
    finally:
        client.kill()


def assert_framing(server: Path) -> None:
    invalid_payloads = [
        b"Content-Length: 2junk\r\n\r\n{}",
        b"Content-Length: 2\r\nContent-Length: 2\r\n\r\n{}",
        b"Content-Length: 0\r\n\r\n",
        b"Content-Length: 16777217\r\n\r\n",
        b"Content-Length: 10\r\n\r\n{}",
        b"Content-Length: 2\n\n{}",
        b"X-Length: 2\r\n\r\n{}",
        b"Content-Length: 2\r\nContent-Type: application/json\r\n\r\n{}",
    ]
    for index, payload in enumerate(invalid_payloads):
        result = run_raw(server, payload)
        if result.returncode != 1:
            raise RuntimeError(f"invalid framing case {index} should return code 1")
        if result.stdout:
            raise RuntimeError(f"invalid framing case {index} should not process its body")

    valid_payload = b"".join([
        content_type_frame({
            "jsonrpc": "2.0",
            "id": 31,
            "method": "initialize",
            "params": {},
        }),
        frame({"jsonrpc": "2.0", "id": 32, "method": "shutdown"}),
        frame({"jsonrpc": "2.0", "method": "exit"}),
    ])
    valid = run_raw(server, valid_payload)
    if valid.returncode != 0:
        raise RuntimeError("canonical Content-Type should be accepted")
    if response_by_id(valid.stdout, 31).get("error"):
        raise RuntimeError("initialize with canonical Content-Type should succeed")
    if response_by_id(valid.stdout, 32).get("result", "missing") is not None:
        raise RuntimeError("continuous framing should reach shutdown")


def assert_diagnostics(server: Path) -> None:
    client = LspClient(server)
    try:
        client.send({
            "jsonrpc": "2.0",
            "id": 41,
            "method": "initialize",
            "params": {},
        })
        if client.read().get("id") != 41:
            raise RuntimeError("missing initialize response")

        uri = "file:///diagnostics.tsl"
        client.send({
            "jsonrpc": "2.0",
            "method": "textDocument/didOpen",
            "params": {
                "textDocument": {
                    "uri": uri,
                    "languageId": "tsl",
                    "version": 1,
                    "text": "function broken(",
                },
            },
        })
        diagnostics = client.read(timeout=1)
        if diagnostics.get("method") != "textDocument/publishDiagnostics":
            raise RuntimeError("didOpen should publish diagnostics")
        params = diagnostics.get("params", {})
        if params.get("uri") != uri or not params.get("diagnostics"):
            raise RuntimeError("syntax diagnostics should identify the opened document")

        client.send({"jsonrpc": "2.0", "id": 42, "method": "shutdown"})
        if client.read().get("id") != 42:
            raise RuntimeError("missing shutdown response")
        client.send({"jsonrpc": "2.0", "method": "exit"})
        if client.close_input() != 0:
            raise RuntimeError("diagnostics sequence should shut down cleanly")
    finally:
        client.kill()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    args = parser.parse_args()

    assert_lifecycle(args.server)
    assert_json_rpc_errors(args.server)
    assert_cancellation_and_failures(args.server)
    assert_framing(args.server)
    assert_diagnostics(args.server)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())