#!/usr/bin/env python3
"""
Tiny fake `pylsp` for tests.

Speaks just enough of the LSP wire protocol (Content-Length framing
plus a handful of methods) to let the editor drive Go-to-Definition,
Find References, and the rest of the request/response surface our e2es
rely on.

Why Python and not bash:
  * Bash 3.2 (the macOS default `/bin/bash`) lacks `mapfile`/`readarray`
    and is finicky about negative substring offsets, so a bash script
    that runs cleanly on Linux can silently die mid-handshake on a
    macOS CI runner — and the only symptom is that the editor waits
    forever for a response that never arrives.
  * Python 3 is universally available on the macOS / Linux / Windows
    images we target, has reliable buffered I/O, and lets us emit
    structured progress lines into the URI log so a hung test has
    actionable diagnostics instead of just "180s timeout".

Behaviour:
  * Logs every received URI (tagged with method) to
    `<FAKE_DEVCONTAINER_STATE>/fake_lsp_uris`.
  * Also logs lifecycle breadcrumbs ("startup", "stdin_eof") to the
    same file so a CI tail of that file shows where we got stuck.
  * Answers `textDocument/definition` and `textDocument/references`
    with a configurable Location. Override per-test by writing the
    desired values into:
      - `<state>/fake_lsp_definition_uri`
      - `<state>/fake_lsp_definition_line`
      - `<state>/fake_lsp_definition_character`
    Defaults: `file:///workspaces/proj/util.py`, line 5, char 0.
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

# Per-test scratch dir. The fake-devcontainer harness sets this; if it
# isn't present we fall back to a tmpdir under TMPDIR so a manual
# invocation still has somewhere to write the log.
STATE_DIR = Path(os.environ.get("FAKE_DEVCONTAINER_STATE")
                 or os.path.join(os.environ.get("TMPDIR", "/tmp"), "fake-lsp-state"))
STATE_DIR.mkdir(parents=True, exist_ok=True)
URI_LOG = STATE_DIR / "fake_lsp_uris"
DEF_URI_FILE = STATE_DIR / "fake_lsp_definition_uri"
DEF_LINE_FILE = STATE_DIR / "fake_lsp_definition_line"
DEF_CHAR_FILE = STATE_DIR / "fake_lsp_definition_character"

DEFAULT_DEF_URI = "file:///workspaces/proj/util.py"
DEFAULT_DEF_LINE = 5
DEFAULT_DEF_CHAR = 0


def log(line: str) -> None:
    """Append a single line to the URI log. The log doubles as a
    progress trace for hung tests — CI can dump it to see how far we
    got."""
    try:
        with URI_LOG.open("a") as f:
            f.write(line + "\n")
    except OSError:
        # Logging must never crash the server. If the state dir
        # vanished mid-test, just keep going.
        pass


def read_def_target() -> tuple[str, int, int]:
    """Resolve the definition target the editor will pin a Location at.
    Reads three optional override files, falling back to defaults."""
    def first_line(path: Path) -> str | None:
        try:
            return path.read_text().splitlines()[0]
        except (OSError, IndexError):
            return None

    uri = first_line(DEF_URI_FILE) or DEFAULT_DEF_URI
    line_text = first_line(DEF_LINE_FILE)
    line = int(line_text) if line_text and line_text.strip().isdigit() else DEFAULT_DEF_LINE
    char_text = first_line(DEF_CHAR_FILE)
    char = int(char_text) if char_text and char_text.strip().isdigit() else DEFAULT_DEF_CHAR
    return uri, line, char


def read_message() -> dict | None:
    """Read one Content-Length-framed JSON-RPC message from stdin.
    Returns the parsed JSON object, or None on clean EOF."""
    headers: dict[str, str] = {}
    while True:
        # `sys.stdin.buffer` for byte-accurate reads — Content-Length
        # counts bytes, not characters, and decode-on-read could
        # mis-frame on multibyte payloads.
        raw = sys.stdin.buffer.readline()
        if raw == b"":
            return None  # EOF before any header → graceful shutdown
        line = raw.decode("ascii", errors="replace").rstrip("\r\n")
        if line == "":
            break
        if ":" in line:
            key, _, value = line.partition(":")
            headers[key.strip()] = value.strip()

    length_str = headers.get("Content-Length")
    if not length_str:
        return None
    try:
        length = int(length_str)
    except ValueError:
        return None
    if length <= 0:
        return {}
    body = sys.stdin.buffer.read(length)
    if len(body) < length:
        # Truncated — treat as EOF so callers can clean up.
        return None
    try:
        return json.loads(body)
    except json.JSONDecodeError:
        return None


def send_message(payload: dict) -> None:
    """Write one Content-Length-framed JSON-RPC message to stdout."""
    body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
    header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
    sys.stdout.buffer.write(header)
    sys.stdout.buffer.write(body)
    sys.stdout.buffer.flush()


def first_uri(message: dict) -> str:
    """Pluck a URI we can use as a wait checkpoint. Production LSP
    messages put it under `textDocument.uri` (notifications/requests),
    `params.uri` (some workspace ops), or `rootUri` (the `initialize`
    handshake — the editor identifies the workspace there). Tests
    wait on log lines like `initialize file://...` so checking each
    of these slots in turn matches the editor's actual traffic."""
    params = message.get("params") or {}
    text_doc = params.get("textDocument") or {}
    return (
        text_doc.get("uri")
        or params.get("uri")
        or params.get("rootUri")
        or ""
    )


def handle(message: dict) -> None:
    method = message.get("method") or ""
    msg_id = message.get("id")
    uri = first_uri(message)
    # Always log a `<method-tail> <uri>` line, even when the URI is
    # empty — the test harness waits on prefix matches like
    # `"initialize "` (with trailing space) and a missing URI must
    # not silently swallow the breadcrumb.
    short_method = method.rsplit("/", 1)[-1] or method
    log(f"{short_method} {uri}")

    if method == "initialize":
        send_message({
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {
                "capabilities": {
                    "textDocumentSync": 1,
                    "definitionProvider": True,
                    "hoverProvider": True,
                    "referencesProvider": True,
                }
            },
        })
        return

    if method == "initialized":
        return  # notification, no response

    if method in ("textDocument/didOpen", "textDocument/didChange",
                  "textDocument/didSave", "textDocument/didClose"):
        return  # already logged above

    if method == "textDocument/definition":
        def_uri, def_line, def_char = read_def_target()
        send_message({
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {
                "uri": def_uri,
                "range": {
                    "start": {"line": def_line, "character": def_char},
                    "end": {"line": def_line, "character": def_char + 1},
                },
            },
        })
        return

    if method == "textDocument/references":
        def_uri, def_line, def_char = read_def_target()
        # Single-element reference list — tests pin the *kind* of URI
        # we emit (a wire URI under the in-container workspace) more
        # than the count.
        send_message({
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": [{
                "uri": def_uri,
                "range": {
                    "start": {"line": def_line, "character": def_char},
                    "end": {"line": def_line, "character": def_char + 1},
                },
            }],
        })
        return

    if method == "textDocument/diagnostic":
        send_message({
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {"kind": "full", "items": []},
        })
        return

    if method == "shutdown":
        send_message({"jsonrpc": "2.0", "id": msg_id, "result": None})
        return

    if method == "exit":
        sys.exit(0)

    # Unknown method — keep the editor happy by responding with `null`
    # if it expects a result. Notifications (no id) just fall through.
    if msg_id is not None:
        send_message({"jsonrpc": "2.0", "id": msg_id, "result": None})


def main() -> None:
    log(f"startup pid={os.getpid()} state={STATE_DIR}")
    while True:
        try:
            message = read_message()
        except KeyboardInterrupt:
            log("interrupt")
            return
        if message is None:
            log("stdin_eof")
            return
        try:
            handle(message)
        except BrokenPipeError:
            log("broken_pipe")
            return
        except Exception as exc:  # noqa: BLE001
            log(f"handler_error {type(exc).__name__}: {exc}")
            # Keep the loop alive so a single bad message doesn't kill
            # the whole session.
            continue


if __name__ == "__main__":
    main()
