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

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, LintState
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

    Reformatter = Callable[[LintState], Optional[Sequence[types.TextEdit]]]
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,
        types.SemanticTokenTypes.Operator.value,
        types.SemanticTokenTypes.TypeParameter.value,
        types.SemanticTokenTypes.Variable.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
    basename_based_lookups: dict[str, C] = dataclasses.field(default_factory=dict)
    path_name_based_lookups: dict[str, C] = dataclasses.field(default_factory=dict)
    default_handler: C | None = 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: (
            types.DidOpenTextDocumentParams
            | types.DidChangeTextDocumentParams
            | types.DocumentDiagnosticParams
        ),
    ) -> list[types.Diagnostic] | None: ...


CLI_DIAGNOSTIC_HANDLERS: dict[str, AsyncLinterImpl] = {}
CLI_FORMAT_FILE_HANDLERS: dict[str, "Reformatter"] = {}


LSP_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()
LSP_FORMAT_FILE_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
TEXT_DOC_INLAY_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
DOCUMENT_LINK_HANDLERS: HandlerDispatchTable[C] = HandlerDispatchTable()
_ALIAS_OF = {}


@dataclasses.dataclass(slots=True, frozen=True)
class BasenameMatchingRule:
    rule_type: Literal["basename", "extension"]
    value: str


@dataclasses.dataclass(slots=True, frozen=True)
class SecondaryLanguage:
    language_id: str
    secondary_lookup: Literal["path-name", "basename"] | None = None


@dataclasses.dataclass(slots=True, frozen=True)
class LanguageDispatchRule:
    primary_language_id: str
    basename: str | None
    path_names: Sequence[str]
    secondary_language_ids: Sequence[SecondaryLanguage]
    is_debsrc_packaging_file: bool

    @classmethod
    def new_rule(
        cls,
        primary_language_id: str,
        basename: str | None,
        path_names: str | Sequence[str],
        secondary_language_ids: Sequence[SecondaryLanguage | str] = (),
    ) -> Self:
        path_names_as_seq: Sequence[str] = (
            (path_names,) if isinstance(path_names, str) else tuple(path_names)
        )
        is_debsrc_packaging_file = any(
            pn.startswith("debian/") for pn in path_names_as_seq
        )
        return LanguageDispatchRule(
            primary_language_id,
            basename,
            path_names_as_seq,
            tuple(
                SecondaryLanguage(l) if isinstance(l, str) else l
                for l in secondary_language_ids
            ),
            is_debsrc_packaging_file,
        )


_STANDARD_HANDLERS = {
    types.TEXT_DOCUMENT_FORMATTING: (
        LSP_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: (
                types.DidOpenTextDocumentParams
                | types.DidChangeTextDocumentParams
                | types.DocumentDiagnosticParams
            ),
        ) -> list[types.Diagnostic] | None:
            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, LSP_DIAGNOSTIC_HANDLERS, _lint_wrapper)
        if file_format.is_debsrc_packaging_file:
            for path_name in file_format.path_names:
                if path_name.startswith("debian/"):
                    CLI_DIAGNOSTIC_HANDLERS[path_name] = func

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


def lsp_cli_reformat_document(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    def _wrapper(func: C) -> C:
        for path_name in file_format.path_names:
            if path_name.startswith("debian/"):
                CLI_FORMAT_FILE_HANDLERS[path_name] = func
        return func

    return _wrapper


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


def lsp_document_link(file_format: LanguageDispatchRule) -> Callable[[C], C]:
    return _registering_wrapper(file_format, DOCUMENT_LINK_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.path_names:
        filename_based_handler = filename_based_dispatch.path_name_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.path_name_based_lookups[filename] = handler

    for secondary_language in file_format.secondary_language_ids:
        secondary_table = handler_dict[secondary_language.language_id]
        if secondary_language.secondary_lookup == "path-name":
            if not file_format.path_names:
                raise AssertionError(
                    f"secondary_lookup=path-name requires the language to have path-names. Please correct definition of {file_format.primary_language_id}"
                )
            for filename in file_format.path_names:
                secondary_handler = secondary_table.path_name_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.path_name_based_lookups[filename] = handler
        elif secondary_language.secondary_lookup == "basename":
            basename = file_format.basename
            if not basename:
                raise AssertionError(
                    f"secondary_lookup=basename requires the language to have a basename. Please correct definition of {file_format.primary_language_id}"
                )
            secondary_handler = secondary_table.basename_based_lookups.get(basename)
            if secondary_handler is not None:
                raise AssertionError(
                    f"There is already a handler for basename {basename} under language ID {secondary_language.language_id}"
                )
            secondary_table.basename_based_lookups[basename] = 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_cli_lsp_features_are_loaded() -> None:
    # These imports are needed to force loading of the LSP files. The relevant registration
    # happens as a side effect of the imports.
    import debputy.lsp.languages as lsp_languages
    from debputy.linting.lint_impl import LINTER_FORMATS

    # Ensure no static analysis tool is temped to optimize out the imports. We need them
    # for the side effect.
    assert lsp_languages
    assert LINTER_FORMATS


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

    feature_list = [
        ("diagnostics (lint)", LSP_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", LSP_FORMAT_FILE_HANDLERS),
        ("document link handler", DOCUMENT_LINK_HANDLERS),
    ]
    print("LSP language IDs and their features:")
    all_ids = sorted({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}")