import collections
import dataclasses
import inspect
import sys
from typing import (
    Callable,
    TypeVar,
    Sequence,
    Union,
    Dict,
    List,
    Optional,
    AsyncIterator,
    Self,
    Generic,
    Protocol,
    TYPE_CHECKING,
)

from debputy.commands.debputy_cmd.context import CommandContext
from debputy.commands.debputy_cmd.output import _output_styling
from debputy.lsp.lsp_self_check import LSP_CHECKS

try:
    from pygls.server import LanguageServer
    from debputy.lsp.debputy_ls import DebputyLanguageServer
except ImportError:
    pass

from debputy.linting.lint_util import AsyncLinterImpl
from debputy.lsp.quickfixes import provide_standard_quickfixes_from_diagnostics_ls
from debputy.lsp.text_util import on_save_trim_end_of_line_whitespace

if TYPE_CHECKING:
    import lsprotocol.types as types
else:
    import debputy.lsprotocol.types as types

C = TypeVar("C", bound=Callable)

SEMANTIC_TOKENS_LEGEND = types.SemanticTokensLegend(
    token_types=[
        types.SemanticTokenTypes.Keyword.value,
        types.SemanticTokenTypes.EnumMember.value,
        types.SemanticTokenTypes.Comment.value,
        types.SemanticTokenTypes.String.value,
        types.SemanticTokenTypes.Macro.value,
    ],
    token_modifiers=[],
)
SEMANTIC_TOKEN_TYPES_IDS = {
    t: idx for idx, t in enumerate(SEMANTIC_TOKENS_LEGEND.token_types)
}

DiagnosticHandler = Callable[
    [
        "DebputyLanguageServer",
        Union["types.DidOpenTextDocumentParams", "types.DidChangeTextDocumentParams"],
    ],
    AsyncIterator[Optional[List[types.Diagnostic]]],
]


@dataclasses.dataclass(slots=True)
class LanguageDispatchTable(Generic[C]):
    language_id: str
    filename_based_lookups: Dict[str, C] = dataclasses.field(default_factory=dict)
    default_handler: Optional[C] = None


class HandlerDispatchTable(Generic[C], Dict[str, LanguageDispatchTable[C]]):
    def __missing__(self, key: str) -> LanguageDispatchTable[C]:
        r = LanguageDispatchTable(key)
        self[key] = r
        return r


class DiagnosticHandlerProtocol(Protocol):
    async def __call__(
        self,
        ls: "DebputyLanguageServer",
        params: Union[
            types.DidOpenTextDocumentParams,
            types.DidChangeTextDocumentParams,
            types.DocumentDiagnosticParams,
        ],
    ) -> Union[List[types.Diagnostic], None]: ...


DIAGNOSTIC_HANDLERS: HandlerDispatchTable[DiagnosticHandlerProtocol] = (
    HandlerDispatchTable[DiagnosticHandlerProtocol]()
)
COMPLETER_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
HOVER_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
CODE_ACTION_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
FOLDING_RANGE_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
SEMANTIC_TOKENS_FULL_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
WILL_SAVE_WAIT_UNTIL_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
FORMAT_FILE_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
TEXT_DOC_INLAY_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
_ALIAS_OF = {}


@dataclasses.dataclass(slots=True, frozen=True)
class SecondaryLanguage:
    language_id: str
    filename_based_lookup: bool = False


@dataclasses.dataclass(slots=True, frozen=True)
class LanguageDispatchRule:
    primary_language_id: str
    filenames: Sequence[str]
    secondary_language_ids: Sequence[SecondaryLanguage]

    @classmethod
    def new_rule(
        cls,
        primary_language_id: str,
        filenames: Union[str, Sequence[str]],
        secondary_language_ids: Sequence[Union[SecondaryLanguage, str]] = (),
    ) -> Self:
        return LanguageDispatchRule(
            primary_language_id,
            (filenames,) if isinstance(filenames, str) else tuple(filenames),
            tuple(
                SecondaryLanguage(l) if isinstance(l, str) else l
                for l in secondary_language_ids
            ),
        )


_STANDARD_HANDLERS = {
    types.TEXT_DOCUMENT_FORMATTING: (
        FORMAT_FILE_HANDLERS,
        on_save_trim_end_of_line_whitespace,
    ),
    types.TEXT_DOCUMENT_CODE_ACTION: (
        CODE_ACTION_HANDLERS,
        provide_standard_quickfixes_from_diagnostics_ls,
    ),
    types.TEXT_DOCUMENT_WILL_SAVE_WAIT_UNTIL: (
        WILL_SAVE_WAIT_UNTIL_HANDLERS,
        on_save_trim_end_of_line_whitespace,
    ),
}


def lint_diagnostics(
    file_format: LanguageDispatchRule,
) -> Callable[[AsyncLinterImpl], AsyncLinterImpl]:

    def _wrapper(func: AsyncLinterImpl) -> AsyncLinterImpl:
        if not inspect.iscoroutinefunction(func):
            raise ValueError("Linters are must be async")

        async def _lint_wrapper(
            ls: "DebputyLanguageServer",
            params: Union[
                types.DidOpenTextDocumentParams,
                types.DidChangeTextDocumentParams,
                types.DocumentDiagnosticParams,
            ],
        ) -> Optional[List[types.Diagnostic]]:
            doc = ls.workspace.get_text_document(params.text_document.uri)
            lint_state = ls.lint_state(doc)
            return await lint_state.run_diagnostics(func)

        _register_handler(file_format, DIAGNOSTIC_HANDLERS, _lint_wrapper)

        return func

    return _wrapper


def lsp_completer(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, COMPLETER_HANDLERS)


def lsp_code_actions(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, CODE_ACTION_HANDLERS)


def lsp_hover(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, HOVER_HANDLERS)


def lsp_text_doc_inlay_hints(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, TEXT_DOC_INLAY_HANDLERS)


def lsp_folding_ranges(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, FOLDING_RANGE_HANDLERS)


def lsp_will_save_wait_until(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, WILL_SAVE_WAIT_UNTIL_HANDLERS)


def lsp_format_document(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, FORMAT_FILE_HANDLERS)


def lsp_semantic_tokens_full(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, SEMANTIC_TOKENS_FULL_HANDLERS)


def lsp_standard_handler(
    file_format: LanguageDispatchRule,
    topic: str,
) -> None:
    res = _STANDARD_HANDLERS.get(topic)
    if res is None:
        raise ValueError(f"No standard handler for {topic}")

    table, handler = res

    _register_handler(file_format, table, handler)


def _registering_wrapper(
    file_formats: LanguageDispatchRule,
    handler_dict: HandlerDispatchTable[C],
) -> Callable[[C], C]:
    def _wrapper(func: C) -> C:
        _register_handler(file_formats, handler_dict, func)
        return func

    return _wrapper


def _register_handler(
    file_format: LanguageDispatchRule,
    handler_dict: HandlerDispatchTable[C],
    handler: C,
) -> None:
    primary_table = handler_dict[file_format.primary_language_id]
    filename_based_dispatch = handler_dict[""]

    if primary_table.default_handler is not None:
        raise AssertionError(
            f"There is already a handler for language ID {file_format.primary_language_id}"
        )

    primary_table.default_handler = handler
    for filename in file_format.filenames:
        filename_based_handler = filename_based_dispatch.filename_based_lookups.get(
            filename
        )
        if filename_based_handler is not None:
            raise AssertionError(f"There is already a handler for filename {filename}")
        filename_based_dispatch.filename_based_lookups[filename] = handler

    for secondary_language in file_format.secondary_language_ids:
        secondary_table = handler_dict[secondary_language.language_id]
        if secondary_language.filename_based_lookup:
            for filename in file_format.filenames:
                secondary_handler = secondary_table.filename_based_lookups.get(filename)
                if secondary_handler is not None:
                    raise AssertionError(
                        f"There is already a handler for filename {filename} under language ID {secondary_language.language_id}"
                    )
                secondary_table.filename_based_lookups[filename] = handler
        elif secondary_table.default_handler is not None:
            raise AssertionError(
                f"There is already a primary handler for language ID {secondary_language.language_id}"
            )
        else:
            secondary_table.default_handler = handler


def ensure_lsp_features_are_loaded() -> None:
    # FIXME: This import is needed to force loading of the LSP files. But it only works
    #  for files with a linter (which currently happens to be all of them, but this is
    #  a bit fragile).
    from debputy.linting.lint_impl import LINTER_FORMATS

    assert LINTER_FORMATS


def describe_lsp_features(context: CommandContext) -> None:
    fo = _output_styling(context.parsed_args, sys.stdout)
    ensure_lsp_features_are_loaded()

    feature_list = [
        ("diagnostics (lint)", DIAGNOSTIC_HANDLERS),
        ("code actions/quickfixes", CODE_ACTION_HANDLERS),
        ("completion suggestions", COMPLETER_HANDLERS),
        ("hover docs", HOVER_HANDLERS),
        ("folding ranges", FOLDING_RANGE_HANDLERS),
        ("semantic tokens", SEMANTIC_TOKENS_FULL_HANDLERS),
        ("on-save handler", WILL_SAVE_WAIT_UNTIL_HANDLERS),
        ("inlay hint (doc)", TEXT_DOC_INLAY_HANDLERS),
        ("format file handler", FORMAT_FILE_HANDLERS),
    ]
    print("LSP language IDs and their features:")
    all_ids = sorted(set(lid for _, t in feature_list for lid in t))
    for lang_id in all_ids:
        if lang_id in _ALIAS_OF:
            continue
        features = [n for n, t in feature_list if lang_id in t]
        print(f" * {lang_id}:")
        for feature in features:
            print(f"   - {feature}")

    aliases = collections.defaultdict(list)
    for lang_id in all_ids:
        main_lang = _ALIAS_OF.get(lang_id)
        if main_lang is None:
            continue
        aliases[main_lang].append(lang_id)

    print()
    print("Aliases:")
    for main_id, aliases in aliases.items():
        print(f" * {main_id}: {', '.join(aliases)}")

    print()
    print("General features:")
    for self_check in LSP_CHECKS:
        is_ok = self_check.test()
        if is_ok:
            print(f" * {self_check.feature}: {fo.colored('enabled', fg='green')}")
        else:
            if self_check.is_mandatory:
                disabled = fo.colored(
                    "missing",
                    fg="red",
                    bg="black",
                    style="bold",
                )
            else:
                disabled = fo.colored(
                    "disabled",
                    fg="yellow",
                    bg="black",
                    style="bold",
                )

            if self_check.how_to_fix:
                print(f" * {self_check.feature}: {disabled}")
                print(f"   - {self_check.how_to_fix}")
            else:
                problem_suffix = f" ({self_check.problem})"
                print(f" * {self_check.feature}: {disabled}{problem_suffix}")
