nx_post_support/server/src/lsp_server.py

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
"""Implementation of tool support over LSP."""

from __future__ import annotations

import json
import os
import pathlib
import re
import sys
import threading

from typing import Any, Optional
import operator
from functools import reduce


# **********************************************************
# Update sys.path before importing any bundled libraries.
# **********************************************************
def update_sys_path(path_to_add: str, strategy: str) -> None:
    """Add given path to `sys.path`."""
    if path_to_add not in sys.path and os.path.isdir(path_to_add):
        if strategy == "useBundled":
            sys.path.insert(0, path_to_add)
        elif strategy == "fromEnvironment":
            sys.path.append(path_to_add)


# Ensure that we can import LSP libraries, and other bundled libraries.
update_sys_path(
    os.fspath(pathlib.Path(__file__).parent.parent / "libs"),
    os.getenv("LS_IMPORT_STRATEGY", "useBundled"),
)

# **********************************************************
# Imports needed for the language server goes below this.
# **********************************************************
# pylint: disable=wrong-import-position,import-error
import lsp_jsonrpc as jsonrpc
import lsprotocol.types as lsp
from pygls import uris, workspace
from common.load_data import standard_items
from tools.semantic_tokens import _Highlighter, TOKEN_TYPES, TokenModifier
from tools.completion_items import completion, remove_existing_items, remove_shared_keys
from tools.inlay_hint import InlayHintGenerator
from tools.file_sourcing import get_all_psc_files, read_psc_file
from tools.signature_help import get_signature_help
from lsp_tclserver import TclLanguageServer


WORKSPACE_SETTINGS = {}
GLOBAL_SETTINGS = {}


MAX_WORKERS = 5
LSP_SERVER = TclLanguageServer(name="NX Postprocessor Support", version="0.0.1", max_workers=MAX_WORKERS)

# **********************************************************
# Tool specific code goes below this.
# **********************************************************


# Delete "Linting features" section if your tool is NOT a linter.
# **********************************************************
# Linting features start here
# **********************************************************

#  See `pylint` implementation for a full featured linter extension:
#  Pylint: https://github.com/microsoft/vscode-pylint/blob/main/bundled/tool


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_OPEN)
def did_open(params: lsp.DidOpenTextDocumentParams) -> None:
    """LSP handler for textDocument/didOpen request."""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    LSP_SERVER.compute_diagnostics(document)
    # Also update custom completion and proc docs for this file
    LSP_SERVER.update_poco_completion_for_file(document)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_SAVE)
def did_save(params: lsp.DidSaveTextDocumentParams) -> None:
    """LSP handler for textDocument/didSave request."""
    _ = LSP_SERVER.workspace.get_text_document(params.text_document.uri)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_CLOSE)
def did_close(_: lsp.DidCloseTextDocumentParams) -> None:
    """LSP handler for textDocument/didClose request."""


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_CHANGE)
def did_change(params: lsp.DidChangeTextDocumentParams) -> None:
    """LSP handler for textDocument/didChange request"""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    LSP_SERVER.compute_diagnostics(document)
    LSP_SERVER.update_poco_completion_for_file(document)


@LSP_SERVER.feature(
    lsp.TEXT_DOCUMENT_DIAGNOSTIC,
    lsp.DiagnosticOptions(
        identifier="pull-diagnostics",
        inter_file_dependencies=False,
        workspace_diagnostics=False,
    ),
)
def document_diagnostic(params: lsp.DocumentDiagnosticParams):
    """Return diagnostics for the requested document"""
    was_cached = True
    if (uri := params.text_document.uri) not in LSP_SERVER.diagnostics:
        was_cached = False
        doc = LSP_SERVER.workspace.get_text_document(uri)
        LSP_SERVER.compute_diagnostics(doc)

    version, diagnostics = LSP_SERVER.diagnostics[uri]
    result_id = f"{uri}@{version}"

    if was_cached and result_id == params.previous_result_id:
        return lsp.UnchangedDocumentDiagnosticReport(result_id)

    return lsp.FullDocumentDiagnosticReport(items=diagnostics, result_id=result_id)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_COMPLETION)
def on_completion(params: lsp.CompletionParams) -> lsp.CompletionList:
    from tools.variable_index import build_variable_index
    from tools.completion_items import BUILTIN_VAR_LABELS

    doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)

    # Base items
    poco = [item for items in LSP_SERVER.poco_completion.values() for item in items]
    base_items = standard_items.tcl_keyword_list + standard_items.nx_procs + standard_items.nx_variables + poco

    # Build variable index from current document
    globals_set, procs_locals, proc_ranges = build_variable_index(doc.source)

    # Always include globals (excluding built-ins)
    dynamic_items = []
    for name in sorted(globals_set):
        if name not in BUILTIN_VAR_LABELS:
            dynamic_items.append(lsp.CompletionItem(label=name, kind=lsp.CompletionItemKind.Variable))

    # Include proc-local variables when cursor is inside that proc
    pos = params.position
    if pos is not None:
        for pr in proc_ranges:
            if pr.start_line <= pos.line <= (pr.end_line or pr.start_line):
                for name in sorted(procs_locals.get(pr.name, set())):
                    # Exclude built-ins and globals to avoid duplication
                    if name not in BUILTIN_VAR_LABELS and name not in globals_set:
                        dynamic_items.append(lsp.CompletionItem(label=name, kind=lsp.CompletionItemKind.Variable))
                break

    # Merge with de-duplication for variables only
    merged: list[lsp.CompletionItem] = []
    seen_var_labels: set[str] = set()
    for it in base_items + dynamic_items:
        if getattr(it, "kind", None) == lsp.CompletionItemKind.Variable:
            if it.label in seen_var_labels:
                continue
            seen_var_labels.add(it.label)
        merged.append(it)

    return lsp.CompletionList(is_incomplete=False, items=merged)


# @LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DOCUMENT_SYMBOL)
# def document_symbols(params: lsp.DocumentSymbolParams):
#     doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
#     ast = LSP_SERVER.parser.parse(doc.source)
#     symbols = LSP_SERVER.extract_tcl_symbols(ast)

#     return symbols


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DOCUMENT_SYMBOL)
def document_symbols(params: lsp.DocumentSymbolParams):
    from tools.document_symbols import build_document_symbols

    doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    return build_document_symbols(doc.source)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_INLAY_HINT)
def inlay_hints(params: lsp.InlayHintParams):
    if not GLOBAL_SETTINGS.get("inlayHint", False):
        return []
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)

    # Reuse cached AST
    tree = LSP_SERVER.get_tree(document)

    # Merge proc signatures across files and traverse once
    merged_signatures = {}
    for sigs in LSP_SERVER.proc_signatures.values():
        merged_signatures.update(sigs)

    generator = InlayHintGenerator(merged_signatures)
    tree.accept(generator, recurse=True)
    return generator.hints


@LSP_SERVER.feature(
    lsp.TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL,
    lsp.SemanticTokensLegend(
        token_types=TOKEN_TYPES,
        token_modifiers=[m.name for m in TokenModifier],
    ),
)
def semantic_tokens(params: lsp.SemanticTokensParams):
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)

    data = []
    plugins = []
    hl = _Highlighter(plugins, LSP_SERVER.poco_completion)

    # Reuse cached AST
    tree = LSP_SERVER.get_tree(document)
    tree.accept(hl, recurse=True)

    tokens = hl.tokens()
    for token in tokens:
        data.extend(
            [
                token.line,
                token.offset,
                token.length,
                TOKEN_TYPES.index(token.tok_type),
                reduce(operator.or_, token.tok_modifiers, 0),
            ]
        )
    return lsp.SemanticTokens(data=data)


@LSP_SERVER.feature(
    lsp.TEXT_DOCUMENT_SIGNATURE_HELP,
    lsp.SignatureHelpOptions(trigger_characters=[" ", "\t", "[", ",", "(", "{", '"', "'"], retrigger_characters=list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_:$\"'{}[]() ,\t")),
)
def signature_help(params: lsp.SignatureHelpParams) -> lsp.SignatureHelp | None:
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    tree = LSP_SERVER.get_tree(document)

    # Merge proc signatures across files
    merged_signatures = {}
    for sigs in LSP_SERVER.proc_signatures.values():
        merged_signatures.update(sigs)

    return get_signature_help(document.source, tree, merged_signatures, params.position)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_HOVER)
def hover(params: lsp.HoverParams) -> lsp.Hover:
    pos = params.position
    document_uri = params.text_document.uri
    document = LSP_SERVER.workspace.get_text_document(document_uri)
    col = params.position.character

    try:
        line = document.lines[pos.line]
    except IndexError:
        return None

    # Do not show hover for proc name in its declaration
    from tools.proc_docs import is_proc_declaration_position

    if is_proc_declaration_position(document.source, pos.line, pos.character):
        return None

    # Identify the token under the cursor
    for m in re.finditer(r"\b\w+\b", line):
        if m.start() <= col <= m.end():
            token = m.group(0)
            break
    else:
        return None

    # 1) If token is a known MOM proc/variable, return built-in hover
    command = token
    data = standard_items.json_data
    all_items = data.get("MOM_procs", []) + data.get("mom_variables", [])

    match = next((item for item in all_items if item["label"] == command), None)
    if match and match.get("kind") == "function":
        label = match.get("label", "")
        parameters = match.get("parameters", [])
        param_lines = "\n".join(f"- `{p['name']}`: {p['desc']}" for p in parameters) or "_None_"
        example_data = match.get("example", [])
        example_md = "\n".join(f"{line}" for line in example_data)
        returns_data = match.get("returns", ["None"])
        returns_md = "\n".join(f"- {line}" for line in returns_data)
        doc_md = f"""\
### 📘 {label}

**Purpose**
{match.get("description", "No description available.")}

**Format**
`{match.get("format", label)}`

**Parameters**
{param_lines}

**Return value**
{returns_md}

**Example**
```tcl
{example_md}"""
        return lsp.Hover(lsp.MarkupContent(kind=lsp.MarkupKind.Markdown, value=doc_md))

    # 2) Otherwise, check if the token is a custom proc and show its preceding doc block
    # Build a merged map of proc -> docs gathered during initialization and updates
    proc_docs: dict[str, str] = {}
    for file_docs in LSP_SERVER.proc_docs.values():
        proc_docs.update(file_docs)

    if token in proc_docs:
        return lsp.Hover(lsp.MarkupContent(kind=lsp.MarkupKind.Markdown, value=proc_docs[token]))

    return None


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DEFINITION)
def goto_definition(params: lsp.DefinitionParams):
    """Provide go-to-definition locations for Tcl procs.

    Strategy:
    - Find the token under the cursor.
    - If it matches a custom proc collected in proc_signatures, locate its declaration
      by searching the current document first, then other indexed files.
    - Return a Location pointing to the proc name in its declaration line.
    """
    doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    pos = params.position
    try:
        line = doc.lines[pos.line]
    except IndexError:
        return None

    # Identify token under cursor
    token = None
    for m in re.finditer(r"\b\w+\b", line):
        if m.start() <= pos.character <= m.end():
            token = m.group(0)
            break
    if not token:
        return None

    # Helper to search a single source text for a proc declaration
    def find_decl_in_source(source_text: str, uri: str) -> Optional[lsp.Location]:
        lines = source_text.split("\n")
        pattern = re.compile(r"^\s*proc\s+" + re.escape(token) + r"\b")
        for i, ln in enumerate(lines):
            m = pattern.match(ln)
            if m:
                start_char = ln.find(token)
                if start_char < 0:
                    start_char = max(m.end() - len(token), 0)
                start = lsp.Position(i, start_char)
                end = lsp.Position(i, start_char + len(token))
                return lsp.Location(uri=uri, range=lsp.Range(start=start, end=end))
        return None

    # 1) Search in current document
    loc = find_decl_in_source(doc.source, doc.uri)
    if loc:
        return loc

    # 2) Search in indexed files from proc_signatures
    # Build list of candidate files that declare this token as a proc
    candidate_files: list[str] = []
    for file_path, procs in LSP_SERVER.proc_signatures.items():
        if token in procs:
            candidate_files.append(file_path)

    for fp in candidate_files:
        uri = pathlib.Path(fp).as_uri()
        # Try to get from workspace if available; else read from disk
        try:
            other_doc = LSP_SERVER.workspace.get_text_document(uri)
            source = other_doc.source
        except Exception:
            try:
                source = pathlib.Path(fp).read_text(encoding="utf-8")
            except Exception:
                continue
        loc = find_decl_in_source(source, uri)
        if loc:
            return loc

    return None


# **********************************************************
# Linting features end here
# **********************************************************


# **********************************************************
# Formatting features start here
# **********************************************************
#  Sample implementations:
#  Black: https://github.com/microsoft/vscode-black-formatter/blob/main/bundled/tool


# **********************************************************
# Formatting features ends here
# **********************************************************
@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_FORMATTING)
def formatting(params: lsp.DocumentFormattingParams) -> list[lsp.TextEdit] | None:
    """LSP handler for textDocument/formatting request."""
    doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)

    source = doc.source
    start = lsp.Position(line=0, character=0)
    last_line = source.rsplit("\n", 1)[-1]
    end = lsp.Position(line=source.count("\n"), character=len(last_line))
    if GLOBAL_SETTINGS.get("formatter", True):
        source = LSP_SERVER.format(doc, params.options)
    return [
        lsp.TextEdit(
            range=lsp.Range(start=start, end=end),
            new_text=source,
        )
    ]


# **********************************************************
# Required Language Server Initialization and Exit handlers.
# **********************************************************
@LSP_SERVER.feature(lsp.WORKSPACE_DID_CHANGE_CONFIGURATION)
def did_change_configuration(_: lsp.DidChangeConfigurationParams):
    """LSP Handler for Config Changes"""


@LSP_SERVER.feature(lsp.INITIALIZE)
def initialize(params: lsp.InitializeParams) -> lsp.InitializeResult:
    """LSP handler for initialize request."""
    log_to_output(f"CWD Server: {os.getcwd()}")

    paths = "\r\n   ".join(sys.path)
    log_to_output(f"sys.path used to run Server:\r\n   {paths}")

    GLOBAL_SETTINGS.update(**params.initialization_options.get("globalSettings", {}))

    settings = params.initialization_options["settings"]
    _update_workspace_settings(settings)
    log_to_output(f"Settings used to run Server:\r\n{json.dumps(settings, indent=4, ensure_ascii=False)}\r\n")
    log_to_output(f"Global settings:\r\n{json.dumps(GLOBAL_SETTINGS, indent=4, ensure_ascii=False)}\r\n")
    semantic_tokens_legend = lsp.SemanticTokensLegend(
        token_types=TOKEN_TYPES,
        token_modifiers=[m.name for m in TokenModifier],
    )
    return lsp.InitializeResult(
        capabilities=lsp.ServerCapabilities(
            document_formatting_provider=GLOBAL_SETTINGS.get("formatter", True),
            semantic_tokens_provider=lsp.SemanticTokensOptions(legend=semantic_tokens_legend, full=True, range=False),
            definition_provider=True,
        )
    )


@LSP_SERVER.feature(lsp.INITIALIZED)
def initialized(_params: lsp.InitializedParams):
    """Kick off background indexing to avoid blocking initialization."""

    def index_workspace():
        try:
            root = LSP_SERVER.workspace.root_path
            log_to_output("Background indexing started...")
            psc_files = get_all_psc_files(pathlib.Path(root))
            for psc_file in psc_files:
                poco_files = read_psc_file(psc_file)
                for sourced_layer in poco_files:
                    completion.reset()
                    try:
                        file_root = pathlib.Path(root).joinpath(sourced_layer.subfolder if sourced_layer.subfolder else "")
                        for tcl_file in sourced_layer.files:
                            filepath = pathlib.Path(file_root).joinpath(f"{tcl_file}.tcl")
                            if not filepath.exists():
                                continue
                            completion.reset()
                            document = LSP_SERVER.workspace.get_text_document(filepath.as_uri())
                            tree = LSP_SERVER.parser.parse(document.source)
                            tree.accept(completion, recurse=True)
                            remove_existing_items(completion.custom_functions, LSP_SERVER.poco_completion)
                            LSP_SERVER.poco_completion[str(filepath)] = completion.custom_functions
                            remove_shared_keys(LSP_SERVER.proc_signatures, completion.proc_signatures)
                            LSP_SERVER.proc_signatures[str(filepath)] = completion.proc_signatures
                            from tools.proc_docs import build_proc_docs

                            LSP_SERVER.proc_docs[str(filepath)] = build_proc_docs(tree, document.source)
                    except Exception as e:
                        log_to_output(f"Fehler beim Parsen von {filepath}: {e}")
            log_to_output("Background indexing completed.")
        except Exception as e:
            log_to_output(f"Background indexing failed: {e}")

    threading.Thread(target=index_workspace, name="nxps-indexer", daemon=True).start()


@LSP_SERVER.feature(lsp.EXIT)
def on_exit(_params: Optional[Any] = None) -> None:
    """Handle clean up on exit."""
    jsonrpc.shutdown_json_rpc()


@LSP_SERVER.feature(lsp.SHUTDOWN)
def on_shutdown(_params: Optional[Any] = None) -> None:
    """Handle clean up on shutdown."""
    jsonrpc.shutdown_json_rpc()


def _get_global_defaults():
    return {
        "path": GLOBAL_SETTINGS.get("path", []),
        "interpreter": GLOBAL_SETTINGS.get("interpreter", [sys.executable]),
        "args": GLOBAL_SETTINGS.get("args", []),
        "importStrategy": GLOBAL_SETTINGS.get("importStrategy", "useBundled"),
        "showNotifications": GLOBAL_SETTINGS.get("showNotifications", "off"),
        "formatter": GLOBAL_SETTINGS.get("formatter", True),
        "inlayHint": GLOBAL_SETTINGS.get("inlayHint", True),
    }


def _update_workspace_settings(settings):
    if not settings:
        key = os.getcwd()
        WORKSPACE_SETTINGS[key] = {
            "cwd": key,
            "workspaceFS": key,
            "workspace": uris.from_fs_path(key),
            **_get_global_defaults(),
        }
        return

    for setting in settings:
        key = uris.to_fs_path(setting["workspace"])
        WORKSPACE_SETTINGS[key] = {
            "cwd": key,
            **setting,
            "workspaceFS": key,
        }


def _get_settings_by_path(file_path: pathlib.Path):
    workspaces = {s["workspaceFS"] for s in WORKSPACE_SETTINGS.values()}

    while file_path != file_path.parent:
        str_file_path = str(file_path)
        if str_file_path in workspaces:
            return WORKSPACE_SETTINGS[str_file_path]
        file_path = file_path.parent

    setting_values = list(WORKSPACE_SETTINGS.values())
    return setting_values[0]


def _get_document_key(document: workspace.Document):
    if WORKSPACE_SETTINGS:
        document_workspace = pathlib.Path(document.path)
        workspaces = {s["workspaceFS"] for s in WORKSPACE_SETTINGS.values()}

        # Find workspace settings for the given file.
        while document_workspace != document_workspace.parent:
            if str(document_workspace) in workspaces:
                return str(document_workspace)
            document_workspace = document_workspace.parent

    return None


def _get_settings_by_document(document: workspace.Document | None):
    if document is None or document.path is None:
        return list(WORKSPACE_SETTINGS.values())[0]

    key = _get_document_key(document)
    if key is None:
        # This is either a non-workspace file or there is no workspace.
        key = os.fspath(pathlib.Path(document.path).parent)
        return {
            "cwd": key,
            "workspaceFS": key,
            "workspace": uris.from_fs_path(key),
            **_get_global_defaults(),
        }

    return WORKSPACE_SETTINGS[str(key)]


# *****************************************************
# Logging and notification.
# *****************************************************
def log_to_output(message: str, msg_type: lsp.MessageType = lsp.MessageType.Log) -> None:
    LSP_SERVER.show_message_log(message, msg_type)


def log_error(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Error)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onError", "onWarning", "always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Error)


def log_warning(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Warning)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onWarning", "always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Warning)


def log_always(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Info)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Info)


# *****************************************************
# Start the server.
# *****************************************************
if __name__ == "__main__":
    LSP_SERVER.start_io()