import inspect
import os.path
from typing import (
    Union,
    Optional,
    TypeVar,
    List,
    TYPE_CHECKING,
    Literal,
)
from collections.abc import Sequence, Callable

from debputy import __version__
from debputy.linting.lint_util import AbortTaskError
from debputy.lsp.lsp_features import (
    LSP_DIAGNOSTIC_HANDLERS,
    COMPLETER_HANDLERS,
    HOVER_HANDLERS,
    SEMANTIC_TOKENS_FULL_HANDLERS,
    CODE_ACTION_HANDLERS,
    SEMANTIC_TOKENS_LEGEND,
    WILL_SAVE_WAIT_UNTIL_HANDLERS,
    LSP_FORMAT_FILE_HANDLERS,
    C,
    TEXT_DOC_INLAY_HANDLERS,
    HandlerDispatchTable,
    DOCUMENT_LINK_HANDLERS,
)
from debputy.util import _info, _trace_log

if TYPE_CHECKING:
    import lsprotocol.types as types

    try:
        from pygls.server import LanguageServer
        from pygls.workspace import TextDocument
    except ImportError:
        pass

    from debputy.lsp.debputy_ls import DebputyLanguageServer

    DEBPUTY_LANGUAGE_SERVER = DebputyLanguageServer("debputy", f"v{__version__}")
else:
    import debputy.lsprotocol.types as types

    try:
        from pygls.server import LanguageServer
        from pygls.workspace import TextDocument
        from debputy.lsp.debputy_ls import DebputyLanguageServer

        DEBPUTY_LANGUAGE_SERVER = DebputyLanguageServer("debputy", f"v{__version__}")
    except ImportError:

        class Mock:

            def feature(self, *args, **kwargs):
                return lambda x: x

        DEBPUTY_LANGUAGE_SERVER = Mock()


P = TypeVar("P")
R = TypeVar("R")
L = TypeVar("L", "LanguageServer", "DebputyLanguageServer")


@DEBPUTY_LANGUAGE_SERVER.feature(types.INITIALIZE)
async def _on_initialize(
    ls: "DebputyLanguageServer",
    params: types.InitializeParams,
) -> None:
    await ls.on_initialize(params)


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_DID_OPEN)
async def _open_document(
    ls: "DebputyLanguageServer",
    params: types.DidChangeTextDocumentParams,
) -> None:
    await _open_or_changed_document(ls, params, "Opened")


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_DID_CHANGE)
async def _changed_document(
    ls: "DebputyLanguageServer",
    params: types.DidChangeTextDocumentParams,
) -> None:
    await _open_or_changed_document(ls, params, "Changed")


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_DID_CLOSE)
def _close_document(
    ls: "DebputyLanguageServer",
    params: types.DidCloseTextDocumentParams,
) -> None:
    ls.close_document(params.text_document.uri)


async def _open_or_changed_document(
    ls: "DebputyLanguageServer",
    params: types.DidOpenTextDocumentParams | types.DidChangeTextDocumentParams,
    event_name: Literal["Opened", "Changed"],
) -> None:
    doc_uri = params.text_document.uri
    doc = ls.workspace.get_text_document(doc_uri)

    await _diagnostics_for(
        ls,
        doc,
        params.text_document.version,
        params,
        event_name,
    )


async def _diagnostics_for(
    ls: "DebputyLanguageServer",
    doc: "TextDocument",
    expected_version: int,
    params: (
        types.DidOpenTextDocumentParams
        | types.DidChangeTextDocumentParams
        | types.DocumentDiagnosticParams
    ),
    event_name: Literal["Opened", "Changed"],
) -> None:
    doc_uri = doc.uri
    id_source, language_id, normalized_filename = ls.determine_language_id(doc)
    log_func = _info if event_name == "Opened" else _trace_log
    handler = _resolve_handler(
        LSP_DIAGNOSTIC_HANDLERS,
        language_id,
        normalized_filename,
        log_func,
    )
    if handler is None:
        log_func(
            f"{event_name} document: {doc.path} ({language_id}, {id_source},"
            f" normalized filename: {normalized_filename}) - no diagnostics handler"
        )
        return
    log_func(
        f"{event_name} document: {doc.path} ({language_id}, {id_source}, normalized filename: {normalized_filename})"
        f" - running diagnostics for doc version {expected_version}"
    )

    try:
        diagnostics = await handler(ls, params)
    except AbortTaskError as e:
        _trace_log(f"Aborted lint task: {e.message}")
        return

    ls.record_diagnostics(doc_uri, expected_version, diagnostics or [], False)


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_COMPLETION)
async def _completions(
    ls: "DebputyLanguageServer",
    params: types.CompletionParams,
) -> types.CompletionList | Sequence[types.CompletionItem] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        COMPLETER_HANDLERS,
        "Complete request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_HOVER)
async def _hover(
    ls: "DebputyLanguageServer",
    params: types.CompletionParams,
) -> types.Hover | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        HOVER_HANDLERS,
        "Hover doc request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_INLAY_HINT)
async def _doc_inlay_hint(
    ls: "DebputyLanguageServer",
    params: types.InlayHintParams,
) -> list[types.InlayHint] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        TEXT_DOC_INLAY_HANDLERS,
        "Inlay hint (doc) request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_CODE_ACTION)
async def _code_actions(
    ls: "DebputyLanguageServer",
    params: types.CodeActionParams,
) -> list[types.Command | types.CodeAction] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        CODE_ACTION_HANDLERS,
        "Code action request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_FOLDING_RANGE)
async def _folding_ranges(
    ls: "DebputyLanguageServer",
    params: types.FoldingRangeParams,
) -> Sequence[types.FoldingRange] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        HOVER_HANDLERS,
        "Folding range request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(
    types.TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL,
    types.SemanticTokensRegistrationOptions(
        SEMANTIC_TOKENS_LEGEND,
        full=True,
    ),
)
async def semantic_tokens_full(
    ls: "DebputyLanguageServer",
    params: types.SemanticTokensParams,
) -> types.SemanticTokens | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        SEMANTIC_TOKENS_FULL_HANDLERS,
        "Semantic tokens request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_WILL_SAVE_WAIT_UNTIL)
async def _will_save_wait_until(
    ls: "DebputyLanguageServer",
    params: types.WillSaveTextDocumentParams,
) -> Sequence[types.TextEdit] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        WILL_SAVE_WAIT_UNTIL_HANDLERS,
        "On-save formatting",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_FORMATTING)
async def _format_document(
    ls: "DebputyLanguageServer",
    params: types.WillSaveTextDocumentParams,
) -> Sequence[types.TextEdit] | None:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        LSP_FORMAT_FILE_HANDLERS,
        "Full document formatting",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_DOCUMENT_LINK)
async def _document_link(
    ls: "DebputyLanguageServer",
    params: types.DocumentLinkParams,
) -> Optional[Sequence[types.DocumentLink]]:
    return await _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        DOCUMENT_LINK_HANDLERS,
        "Document links request",
    )


async def _dispatch_standard_handler(
    ls: "DebputyLanguageServer",
    doc_uri: str,
    params: P,
    handler_table: HandlerDispatchTable[C],
    request_type: str,
) -> R | None:
    doc = ls.workspace.get_text_document(doc_uri)

    id_source, language_id, normalized_filename = ls.determine_language_id(doc)
    handler = _resolve_handler(handler_table, language_id, normalized_filename, _info)
    if handler is None:
        _info(
            f"{request_type} for document: {doc.path} ({language_id}, {id_source},"
            f" normalized filename: {normalized_filename}) - no handler"
        )
        return None
    _info(
        f"{request_type} for document: {doc.path} ({language_id}, {id_source},"
        f" normalized filename: {normalized_filename}) - delegating to handler {handler.__qualname__}"
    )

    if inspect.iscoroutinefunction(handler):
        return await handler(ls, params)
    return handler(ls, params)


def _resolve_handler(
    handler_table: HandlerDispatchTable[C],
    language_id: str,
    normalized_filename: str,
    log_func: Callable[[str], None],
) -> C | None:
    dispatch_table = handler_table.get(language_id) if language_id != "" else None
    log_func(f"resolve_handler({language_id=}, {normalized_filename=})")
    if dispatch_table is None:
        filename_based_table = handler_table[""]
        m = filename_based_table.basename_based_lookups.get(
            os.path.basename(normalized_filename)
        )
        if m is not None:
            return m
        return filename_based_table.path_name_based_lookups.get(normalized_filename)
    m = dispatch_table.basename_based_lookups.get(
        os.path.basename(normalized_filename),
        dispatch_table.default_handler,
    )
    if m is not None:
        return m
    return dispatch_table.path_name_based_lookups.get(
        normalized_filename, dispatch_table.default_handler
    )