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

from __future__ import annotations

import copy
import json
import os
import pathlib
import re
import sys
import sysconfig
import traceback
from typing import Any, Optional, Sequence, Tuple
import re


# **********************************************************
# 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 lsp_utils as utils
import lsprotocol.types as lsp
from pygls import server, uris, workspace
from pygls.workspace.text_document import TextDocument
from common.load_data import standard_items
from common.formatter import format_tcl
from tclint.parser import Parser
from tclint.format import Formatter, FormatterOpts
from tools.semantic_tokens import (
    SemanticTokenCollector,
    collect_semantic_tokens,
    encode_tokens,
)
from tools.nx_plugins import commands


class TclLanguageServer(server.LanguageServer):
    def format(
        self,
        document: TextDocument,
        options: lsp.FormattingOptions,
        range: Optional[Tuple[int, int]] = None,
    ):
        parser = Parser()
        parser._commands.update(commands)
        tree = parser.parse(document.source)

        log_to_output(tree.pretty(2))

        indent = "\t" if not options.insert_spaces else " " * options.tab_size

        formatter = Formatter(
            FormatterOpts(
                indent=indent,
                spaces_in_braces=False,
                max_blank_lines=500,
                indent_namespace_eval=True,
            ),
        )

        if range is not None:
            start, end = range
            return formatter.format_partial(document.source[start:end], parser)

        return formatter.format_top(document.source, parser)


WORKSPACE_SETTINGS = {}
GLOBAL_SETTINGS = {}
RUNNER = pathlib.Path(__file__).parent / "lsp_runner.py"

TOKEN_TYPES = [
    "command",
    "variable",
    "function",
    "string",
    "number",
    "keyword",
    "comment",
]
TOKEN_MODIFIERS = []

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

# **********************************************************
# Tool specific code goes below this.
# **********************************************************
TOOL_MODULE = "nx-post-support"
TOOL_DISPLAY = "NX Postprocessor Support"
TOOL_ARGS = []  # default arguments always passed to your tool.

# 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.feature(lsp.TEXT_DOCUMENT_DID_SAVE)
def did_save(params: lsp.DidSaveTextDocumentParams) -> None:
    """LSP handler for textDocument/didSave request."""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_CLOSE)
def did_close(params: 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"""
    log_to_output("Document has changed")


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_COMPLETION)
def on_completion(params: lsp.CompletionParams) -> list[lsp.CompletionItem]:
    items = (
        standard_items.tcl_keyword_list
        + standard_items.nx_procs
        + standard_items.nx_variables
    )
    return lsp.CompletionList(is_incomplete=False, items=items)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL)
def on_semantic_tokens(params: lsp.SemanticTokensParams):
    doc = LSP_SERVER.workspace.get_document(params.text_document.uri)
    code = doc.source

    tokens = collect_semantic_tokens(code)
    data = encode_tokens(tokens)

    return lsp.SemanticTokens(data=data)


@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

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

    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 not match:
        return None

    if not match.get("kind") == "function":
        return None

    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))


# **********************************************************
# 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))

    formatted = LSP_SERVER.format(doc, params.options)
    return [
        lsp.TextEdit(
            range=lsp.Range(start=start, end=end),
            new_text=formatted,
        )
    ]


# **********************************************************
# Required Language Server Initialization and Exit handlers.
# **********************************************************
@LSP_SERVER.feature(lsp.WORKSPACE_DID_CHANGE_CONFIGURATION)
def did_change_configuration(params: lsp.DidChangeConfigurationParams):
    settings = params.settings

    log_to_output(str(settings))


@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=TOKEN_MODIFIERS,
    )
    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
            ),
        )
    )


@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),
    }


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()