# Copyright 2026 Marimo. All rights reserved.

"""LSP handlers for marimo."""

from __future__ import annotations

import atexit
import importlib.metadata
import typing

import lsprotocol.types as lsp
import msgspec
from marimo._convert.converters import MarimoConvert
from pygls.lsp.server import LanguageServer
from pygls.uris import to_fs_path, uri_scheme

from marimo_lsp.api import handle_api_command
from marimo_lsp.completions import get_completions
from marimo_lsp.diagnostics import GraphUpdaterRegistry
from marimo_lsp.loggers import get_logger
from marimo_lsp.models import ApiRequest, ConvertRequest
from marimo_lsp.sessions import Sessions

logger = get_logger()

if typing.TYPE_CHECKING:
    from marimo_lsp.kernels import Kernels


def create_server(*, kernels: Kernels) -> LanguageServer:  # noqa: C901, PLR0915
    """Create the marimo LSP server."""
    server = LanguageServer(
        name="marimo-lsp",
        version=importlib.metadata.version("marimo-lsp"),
        notebook_document_sync=lsp.NotebookDocumentSyncOptions(
            notebook_selector=[
                lsp.NotebookDocumentFilterWithCells(
                    notebook="marimo-notebook",
                    cells=[
                        lsp.NotebookCellLanguage(language="sql"),
                        lsp.NotebookCellLanguage(language="python"),
                        lsp.NotebookCellLanguage(language="mo-python"),
                        lsp.NotebookCellLanguage(language="markdown"),
                    ],
                ),
            ],
            save=True,
        ),
    )
    sessions = Sessions(server, kernels=kernels)
    graph_registry = GraphUpdaterRegistry(server)

    # Register atexit handler to ensure kernel processes are cleaned up
    # when the LSP server exits (e.g., extension host restart, VS Code close).
    # This prevents orphaned kernel subprocesses from consuming memory.
    atexit.register(sessions.close_all)

    # Lsp Features
    @server.feature(lsp.SHUTDOWN)
    def shutdown(params: None) -> None:
        del params
        sessions.close_all()

    @server.feature(lsp.NOTEBOOK_DOCUMENT_DID_OPEN)
    async def did_open(params: lsp.DidOpenNotebookDocumentParams) -> None:
        logger.info(f"notebookDocument/didOpen {params.notebook_document.uri}")
        sessions.attach(params.notebook_document.uri, server.workspace)

        # Immediate compile + publish — needed for initial cell ordering
        updater = graph_registry.get_or_create(params.notebook_document.uri)
        updater.flush()

    @server.feature(lsp.NOTEBOOK_DOCUMENT_DID_CHANGE)
    async def did_change(params: lsp.DidChangeNotebookDocumentParams) -> None:
        logger.info(f"notebookDocument/didChange {params.notebook_document.uri}")
        sessions.sync(params.notebook_document.uri, server.workspace)
        logger.info(f"Synced session {params.notebook_document.uri}")

        # Schedule debounced recompilation — compiles after 150ms of quiet
        updater = graph_registry.get_or_create(params.notebook_document.uri)
        updater.schedule()

    @server.feature(lsp.NOTEBOOK_DOCUMENT_DID_SAVE)
    async def did_save(params: lsp.DidSaveNotebookDocumentParams) -> None:
        logger.info(f"notebookDocument/didSave {params.notebook_document.uri}")
        sessions.sync(params.notebook_document.uri, server.workspace)
        logger.info(f"Synced session {params.notebook_document.uri}")

    @server.feature(lsp.NOTEBOOK_DOCUMENT_DID_CLOSE)
    async def did_close(params: lsp.DidCloseNotebookDocumentParams) -> None:
        logger.info(f"notebookDocument/didClose {params.notebook_document.uri}")

        # Clean up graph updater (cancels pending debounce timer)
        graph_registry.remove(params.notebook_document.uri)

        # Only close untitled sessions, when closing documents...
        # Others can stay alive
        if params.notebook_document.uri.startswith("untitled:"):
            sessions.close(params.notebook_document.uri)
            logger.info(f"Closed {params.notebook_document.uri}")
        else:
            sessions.detach(params.notebook_document.uri)

    @server.feature(lsp.TEXT_DOCUMENT_DIAGNOSTIC)
    def diagnostics(params: lsp.DocumentDiagnosticParams):
        """Provide diagnostics for marimo notebooks.

        The ``textDocument/diagnostic`` request is pull-based.  We flush any
        pending debounced recompilation so that the response reflects the
        latest cell state.
        """
        logger.info(f"textDocument/diagnostic {params.text_document.uri}")

        notebook = server.workspace.get_notebook_document(
            cell_uri=params.text_document.uri
        )

        if not notebook:
            logger.debug("No target notebook found for diagnostics")
            return lsp.RelatedFullDocumentDiagnosticReport(kind="full", items=[])

        # Ensure the graph is up-to-date before responding
        updater = graph_registry.get_or_create(notebook.uri)
        updater.flush()

        # Diagnostics are pushed via publishDiagnostics notifications
        # from _recompile(); the pull response is empty.
        return lsp.RelatedFullDocumentDiagnosticReport(kind="full", items=[])

    @server.feature(
        lsp.TEXT_DOCUMENT_CODE_ACTION,
        lsp.CodeActionOptions(
            code_action_kinds=[lsp.CodeActionKind.RefactorRewrite],
            resolve_provider=False,
        ),
    )
    def code_actions(params: lsp.CodeActionParams):
        """Provide code actions for Python files to convert to marimo."""
        logger.info(f"textDocument/codeAction {params.text_document.uri}")

        scheme = uri_scheme(params.text_document.uri)
        if scheme and scheme.endswith("notebook-cell"):
            # No code actions for notebook cells (for now)
            return []

        actions: list[lsp.CodeAction] = []
        filename = to_fs_path(params.text_document.uri)
        if filename and filename.endswith((".py", ".ipynb")):
            actions.append(
                lsp.CodeAction(
                    title="Convert to marimo notebook",
                    kind=lsp.CodeActionKind.RefactorRewrite,
                    command=lsp.Command(
                        title="Convert to marimo notebook",
                        command="marimo.convert",
                        arguments=[{"uri": params.text_document.uri}],
                    ),
                )
            )

        return actions

    @server.feature(
        lsp.TEXT_DOCUMENT_COMPLETION,
        lsp.CompletionOptions(
            trigger_characters=["@"],
            resolve_provider=False,
        ),
    )
    def completions(ls: LanguageServer, params: lsp.CompletionParams):
        """Provide completions for marimo cells."""
        logger.info(f"textDocument/completion {params.text_document.uri}")

        scheme = uri_scheme(params.text_document.uri)
        if scheme and scheme.endswith("notebook-cell"):
            # No completions for notebook cells (for now)
            return []

        return get_completions(ls, params)

    @server.command("marimo.api")
    async def api(ls: LanguageServer, params: typing.Any):  # noqa: ANN401
        """Unified API endpoint for all marimo internal methods."""
        logger.info("marimo.api")
        args = msgspec.convert(params, type=ApiRequest)
        return await handle_api_command(ls, sessions, args.method, args.params)

    @server.command("marimo.convert")
    async def convert(ls: LanguageServer, params: typing.Any):  # noqa: ANN401
        """Convert a Python file to marimo format and create a new file."""
        logger.info("marimo.convert")

        args = msgspec.convert(params, type=ConvertRequest)
        text_document = ls.workspace.get_text_document(args.uri)
        filename = text_document.filename

        if filename is None:
            return

        if filename.endswith(".ipynb"):
            ir = MarimoConvert.from_ipynb(text_document.source)
            new_filename = filename.replace(".ipynb", "_mo.py")
        else:
            ir = MarimoConvert.from_non_marimo_python_script(text_document.source)
            new_filename = filename.replace(".py", "_mo.py")

        new_uri = text_document.uri.replace(filename, new_filename)
        new_text = ir.to_py()
        result = await ls.workspace_apply_edit_async(
            lsp.ApplyWorkspaceEditParams(
                label=f"converted {filename} → {new_filename}",
                edit=lsp.WorkspaceEdit(
                    document_changes=[
                        lsp.CreateFile(
                            kind="create",
                            uri=new_uri,
                            options=lsp.CreateFileOptions(
                                overwrite=False,
                                ignore_if_exists=True,
                            ),
                        ),
                        lsp.TextDocumentEdit(
                            text_document=lsp.OptionalVersionedTextDocumentIdentifier(
                                uri=new_uri,
                                version=None,
                            ),
                            edits=[
                                lsp.TextEdit(
                                    new_text=new_text,
                                    range=lsp.Range(
                                        start=lsp.Position(line=0, character=0),
                                        end=lsp.Position(line=0, character=0),
                                    ),
                                )
                            ],
                        ),
                    ],
                ),
            )
        )
        if result.applied:
            await ls.window_show_document_async(
                lsp.ShowDocumentParams(
                    uri=new_uri,
                    external=False,
                    take_focus=True,
                    selection=None,
                )
            )

    logger.info("All handlers registered successfully")

    return server
