layonara-forge/packages/frontend/src/lib/lspClient.ts

import { toSocket } from "vscode-ws-jsonrpc";
import { createWebSocketConnection } from "vscode-ws-jsonrpc";
import type * as lsp from "vscode-languageserver-protocol";
import type { editor, languages, IDisposable, Uri, MarkerSeverity } from "monaco-editor";

const LSP_TRACE = false;

function trace(...args: unknown[]) {
  if (LSP_TRACE) console.debug("[LSP]", ...args);
}

export type LspStatus = "disconnected" | "connecting" | "initializing" | "ready" | "error";

export class LspClient {
  private ws: WebSocket | null = null;
  private connection: ReturnType<typeof createWebSocketConnection> | null = null;
  private disposables: IDisposable[] = [];
  private documentVersions = new Map<string, number>();
  private statusListeners = new Set<(status: LspStatus) => void>();
  private _status: LspStatus = "disconnected";
  private monacoInstance: typeof import("monaco-editor") | null = null;

  get status() {
    return this._status;
  }

  private setStatus(s: LspStatus) {
    this._status = s;
    for (const cb of this.statusListeners) cb(s);
  }

  onStatusChange(cb: (status: LspStatus) => void): () => void {
    this.statusListeners.add(cb);
    return () => this.statusListeners.delete(cb);
  }

  async connect(monaco: typeof import("monaco-editor")): Promise<void> {
    this.monacoInstance = monaco;
    this.setStatus("connecting");

    const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
    const url = `${protocol}//${window.location.host}/ws/lsp`;

    return new Promise<void>((resolve, reject) => {
      const ws = new WebSocket(url);
      this.ws = ws;

      ws.onopen = async () => {
        try {
          const socket = toSocket(ws);
          this.connection = createWebSocketConnection(socket, {
            error: (msg) => console.error("[LSP]", msg),
            warn: (msg) => console.warn("[LSP]", msg),
            info: (msg) => trace(msg),
            log: (msg) => trace(msg),
          });

          this.setupNotificationHandlers(monaco);
          this.connection.listen();

          this.setStatus("initializing");
          await this.initialize();
          this.registerProviders(monaco);
          this.setStatus("ready");
          resolve();
        } catch (err) {
          this.setStatus("error");
          reject(err);
        }
      };

      ws.onerror = () => {
        this.setStatus("error");
        reject(new Error("WebSocket connection failed"));
      };

      ws.onclose = () => {
        this.setStatus("disconnected");
      };
    });
  }

  private async initialize(): Promise<void> {
    if (!this.connection) throw new Error("No connection");

    const initParams: lsp.InitializeParams = {
      processId: null,
      capabilities: {
        textDocument: {
          completion: {
            completionItem: {
              snippetSupport: false,
              documentationFormat: ["plaintext", "markdown"],
            },
          },
          hover: {
            contentFormat: ["plaintext", "markdown"],
          },
          synchronization: {
            didSave: true,
          },
          publishDiagnostics: {
            relatedInformation: true,
          },
        },
      },
      rootUri: "file:///workspace",
      workspaceFolders: null,
    };

    trace("→ initialize", initParams);
    const result = await this.connection.sendRequest("initialize", initParams);
    trace("← initialize", result);

    this.connection.sendNotification("initialized", {});
    trace("→ initialized");
  }

  private setupNotificationHandlers(monaco: typeof import("monaco-editor")): void {
    if (!this.connection) return;

    this.connection.onNotification(
      "textDocument/publishDiagnostics",
      (params: lsp.PublishDiagnosticsParams) => {
        trace("← publishDiagnostics", params);
        const uri = monaco.Uri.parse(params.uri);
        const model = monaco.editor.getModel(uri);
        if (!model) return;

        const markers: editor.IMarkerData[] = params.diagnostics.map((d) => ({
          severity: this.mapSeverity(monaco, d.severity),
          startLineNumber: d.range.start.line + 1,
          startColumn: d.range.start.character + 1,
          endLineNumber: d.range.end.line + 1,
          endColumn: d.range.end.character + 1,
          message: d.message,
          source: d.source,
        }));

        monaco.editor.setModelMarkers(model, "nwscript-lsp", markers);
      },
    );
  }

  private mapSeverity(
    monaco: typeof import("monaco-editor"),
    severity: lsp.DiagnosticSeverity | undefined,
  ): MarkerSeverity {
    switch (severity) {
      case 1:
        return monaco.MarkerSeverity.Error;
      case 2:
        return monaco.MarkerSeverity.Warning;
      case 3:
        return monaco.MarkerSeverity.Info;
      case 4:
        return monaco.MarkerSeverity.Hint;
      default:
        return monaco.MarkerSeverity.Info;
    }
  }

  private registerProviders(monaco: typeof import("monaco-editor")): void {
    this.disposables.push(
      monaco.languages.registerCompletionItemProvider("nwscript", {
        triggerCharacters: ["."],
        provideCompletionItems: async (model, position) => {
          if (!this.connection) return { suggestions: [] };

          const params: lsp.CompletionParams = {
            textDocument: { uri: model.uri.toString() },
            position: {
              line: position.lineNumber - 1,
              character: position.column - 1,
            },
          };

          trace("→ completion", params);
          const result: lsp.CompletionItem[] | lsp.CompletionList | null =
            await this.connection.sendRequest("textDocument/completion", params);
          trace("← completion", result);

          if (!result) return { suggestions: [] };
          const items = Array.isArray(result) ? result : result.items;

          const suggestions: languages.CompletionItem[] = items.map((item) => ({
            label: item.label,
            kind: this.mapCompletionKind(monaco, item.kind),
            insertText: item.insertText ?? item.label,
            detail: item.detail,
            documentation: item.documentation
              ? typeof item.documentation === "string"
                ? item.documentation
                : { value: item.documentation.value }
              : undefined,
            range: {
              startLineNumber: position.lineNumber,
              startColumn: position.column,
              endLineNumber: position.lineNumber,
              endColumn: position.column,
            },
          }));

          return { suggestions };
        },
      }),
    );

    this.disposables.push(
      monaco.languages.registerHoverProvider("nwscript", {
        provideHover: async (model, position) => {
          if (!this.connection) return null;

          const params: lsp.HoverParams = {
            textDocument: { uri: model.uri.toString() },
            position: {
              line: position.lineNumber - 1,
              character: position.column - 1,
            },
          };

          trace("→ hover", params);
          const result: lsp.Hover | null = await this.connection.sendRequest(
            "textDocument/hover",
            params,
          );
          trace("← hover", result);

          if (!result?.contents) return null;

          const contents = Array.isArray(result.contents)
            ? result.contents.map((c) =>
                typeof c === "string" ? { value: c } : { value: c.value },
              )
            : typeof result.contents === "string"
              ? [{ value: result.contents }]
              : "value" in result.contents
                ? [{ value: result.contents.value }]
                : [{ value: String(result.contents) }];

          return {
            contents,
            range: result.range
              ? {
                  startLineNumber: result.range.start.line + 1,
                  startColumn: result.range.start.character + 1,
                  endLineNumber: result.range.end.line + 1,
                  endColumn: result.range.end.character + 1,
                }
              : undefined,
          };
        },
      }),
    );
  }

  private mapCompletionKind(
    monaco: typeof import("monaco-editor"),
    kind: lsp.CompletionItemKind | undefined,
  ): languages.CompletionItemKind {
    const k = monaco.languages.CompletionItemKind;
    switch (kind) {
      case 1:
        return k.Text;
      case 2:
        return k.Method;
      case 3:
        return k.Function;
      case 4:
        return k.Constructor;
      case 5:
        return k.Field;
      case 6:
        return k.Variable;
      case 7:
        return k.Class;
      case 8:
        return k.Interface;
      case 9:
        return k.Module;
      case 10:
        return k.Property;
      case 13:
        return k.Enum;
      case 14:
        return k.Keyword;
      case 15:
        return k.Snippet;
      case 21:
        return k.Constant;
      case 22:
        return k.Struct;
      default:
        return k.Text;
    }
  }

  notifyDidOpen(uri: Uri, languageId: string, text: string): void {
    if (!this.connection || this._status !== "ready") return;

    this.documentVersions.set(uri.toString(), 1);
    const params: lsp.DidOpenTextDocumentParams = {
      textDocument: {
        uri: uri.toString(),
        languageId,
        version: 1,
        text,
      },
    };
    trace("→ didOpen", params.textDocument.uri);
    this.connection.sendNotification("textDocument/didOpen", params);
  }

  notifyDidChange(uri: Uri, text: string): void {
    if (!this.connection || this._status !== "ready") return;

    const version = (this.documentVersions.get(uri.toString()) ?? 0) + 1;
    this.documentVersions.set(uri.toString(), version);
    const params: lsp.DidChangeTextDocumentParams = {
      textDocument: { uri: uri.toString(), version },
      contentChanges: [{ text }],
    };
    trace("→ didChange", params.textDocument.uri);
    this.connection.sendNotification("textDocument/didChange", params);
  }

  notifyDidClose(uri: Uri): void {
    if (!this.connection || this._status !== "ready") return;

    this.documentVersions.delete(uri.toString());
    const params: lsp.DidCloseTextDocumentParams = {
      textDocument: { uri: uri.toString() },
    };
    trace("→ didClose", params.textDocument.uri);
    this.connection.sendNotification("textDocument/didClose", params);
  }

  dispose(): void {
    for (const d of this.disposables) d.dispose();
    this.disposables = [];
    this.connection?.dispose();
    this.connection = null;
    if (this.ws && this.ws.readyState !== WebSocket.CLOSED) {
      this.ws.close();
    }
    this.ws = null;
    this.documentVersions.clear();
    this.setStatus("disconnected");
  }
}

let globalClient: LspClient | null = null;

export function getLspClient(): LspClient {
  if (!globalClient) {
    globalClient = new LspClient();
  }
  return globalClient;
}