import asyncio
from typing import (
    Dict,
    Sequence,
    Union,
    Optional,
    TypeVar,
    List,
    TYPE_CHECKING,
    Callable,
)

from debputy import __version__
from debputy.lsp.lsp_features import (
    DIAGNOSTIC_HANDLERS,
    COMPLETER_HANDLERS,
    HOVER_HANDLERS,
    SEMANTIC_TOKENS_FULL_HANDLERS,
    CODE_ACTION_HANDLERS,
    SEMANTIC_TOKENS_LEGEND,
    WILL_SAVE_WAIT_UNTIL_HANDLERS,
    FORMAT_FILE_HANDLERS,
    C,
    TEXT_DOC_INLAY_HANDLERS,
    HandlerDispatchTable,
)
from debputy.util import _info

_DOCUMENT_VERSION_TABLE: Dict[str, int] = {}


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


def is_doc_at_version(uri: str, version: int) -> bool:
    dv = _DOCUMENT_VERSION_TABLE.get(uri)
    return dv == version


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


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


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

    def _diag_publisher(
        uri: str,
        diagnostics: Optional[List[types.Diagnostic]],
        version: int,
        _is_partial: bool,
    ) -> None:
        ls.publish_diagnostics(uri, diagnostics, version)

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


async def _diagnostics_for(
    ls: "DebputyLanguageServer",
    doc: "TextDocument",
    expected_version: int,
    params: Union[
        types.DidOpenTextDocumentParams,
        types.DidChangeTextDocumentParams,
        types.DocumentDiagnosticParams,
    ],
    publisher: Callable[[str, Optional[List["types.Diagnostic"]], int, bool], None],
) -> None:
    doc_uri = doc.uri
    _DOCUMENT_VERSION_TABLE[doc_uri] = expected_version
    id_source, language_id, normalized_filename = ls.determine_language_id(doc)
    handler = _resolve_handler(DIAGNOSTIC_HANDLERS, language_id, normalized_filename)
    if handler is None:
        _info(
            f"Opened/Changed document: {doc.path} ({language_id}, {id_source},"
            f" normalized filename: {normalized_filename}) - no diagnostics handler"
        )
        return
    _info(
        f"Opened/Changed document: {doc.path} ({language_id}, {id_source}, normalized filename: {normalized_filename})"
        f" - running diagnostics for doc version {expected_version}"
    )
    last_publish_count = -1

    diagnostics_scanner = handler(ls, params)
    diagnostics: Optional[List[types.Diagnostic]] = None
    async for diagnostics in diagnostics_scanner:
        await asyncio.sleep(0)
        if not is_doc_at_version(doc_uri, expected_version):
            # This basically happens with very edit, so lets not notify the client
            # for that.
            _info(
                f"Cancel (obsolete) diagnostics for doc version {expected_version}: document version changed"
            )
            break
        if diagnostics is None or last_publish_count != len(diagnostics):
            last_publish_count = len(diagnostics) if diagnostics is not None else 0
            publisher(
                doc.uri,
                diagnostics,
                expected_version,
                True,
            )
    publisher(
        doc.uri,
        diagnostics,
        expected_version,
        False,
    )


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


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_HOVER)
def _hover(
    ls: "DebputyLanguageServer",
    params: types.CompletionParams,
) -> Optional[types.Hover]:
    return _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        HOVER_HANDLERS,
        "Hover doc request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_INLAY_HINT)
def _doc_inlay_hint(
    ls: "DebputyLanguageServer",
    params: types.InlayHintParams,
) -> Optional[List[types.InlayHint]]:
    return _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)
def _code_actions(
    ls: "DebputyLanguageServer",
    params: types.CodeActionParams,
) -> Optional[List[Union[types.Command, types.CodeAction]]]:
    return _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        CODE_ACTION_HANDLERS,
        "Code action request",
    )


@DEBPUTY_LANGUAGE_SERVER.feature(types.TEXT_DOCUMENT_FOLDING_RANGE)
def _folding_ranges(
    ls: "DebputyLanguageServer",
    params: types.FoldingRangeParams,
) -> Optional[Sequence[types.FoldingRange]]:
    return _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,
    ),
)
def _semantic_tokens_full(
    ls: "DebputyLanguageServer",
    params: types.SemanticTokensParams,
) -> Optional[types.SemanticTokens]:
    return _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)
def _will_save_wait_until(
    ls: "DebputyLanguageServer",
    params: types.WillSaveTextDocumentParams,
) -> Optional[Sequence[types.TextEdit]]:
    return _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)
def _format_document(
    ls: "DebputyLanguageServer",
    params: types.WillSaveTextDocumentParams,
) -> Optional[Sequence[types.TextEdit]]:
    return _dispatch_standard_handler(
        ls,
        params.text_document.uri,
        params,
        FORMAT_FILE_HANDLERS,
        "Full document formatting",
    )


def _dispatch_standard_handler(
    ls: "DebputyLanguageServer",
    doc_uri: str,
    params: P,
    handler_table: HandlerDispatchTable[C],
    request_type: str,
) -> Optional[R]:
    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)
    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"
    )

    return handler(
        ls,
        params,
    )


def _resolve_handler(
    handler_table: HandlerDispatchTable[C],
    language_id: str,
    normalized_filename: str,
) -> Optional[C]:
    dispatch_table = handler_table.get(language_id) if language_id != "" else None
    _info(f"LID: {language_id} - {dispatch_table}")
    if dispatch_table is None:
        filename_based_table = handler_table[""]
        return filename_based_table.filename_based_lookups.get(normalized_filename)
    return dispatch_table.filename_based_lookups.get(
        normalized_filename, dispatch_table.default_handler
    )
