from typing import (
    Literal,
    TypedDict,
    Union,
    TypeVar,
    Dict,
    Optional,
    List,
    cast,
    NotRequired,
    TYPE_CHECKING,
)
from collections.abc import Callable, Iterable, Mapping

from debputy.lsp.diagnostics import DiagnosticData
from debputy.util import _warn

if TYPE_CHECKING:
    import lsprotocol.types as types
    from debputy.linting.lint_util import LintState
else:
    import debputy.lsprotocol.types as types


try:
    from debputy.lsp.vendoring._deb822_repro.locatable import (
        Position as TEPosition,
        Range as TERange,
    )

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


CodeActionName = Literal[
    "correct-text",
    "remove-line",
    "remove-range",
    "insert-text-on-line-after-diagnostic",
]


class CorrectTextCodeAction(TypedDict):
    code_action: Literal["correct-text"]
    correct_value: str
    proposed_title: NotRequired[str]


class InsertTextOnLineAfterDiagnosticCodeAction(TypedDict):
    code_action: Literal["insert-text-on-line-after-diagnostic"]
    text_to_insert: str


class RemoveLineCodeAction(TypedDict):
    code_action: Literal["remove-line"]


class RemoveRangeCodeAction(TypedDict):
    code_action: Literal["remove-range"]
    proposed_title: NotRequired[str]


def propose_correct_text_quick_fix(
    correct_value: str,
    *,
    proposed_title: str | None = None,
) -> CorrectTextCodeAction:
    r: CorrectTextCodeAction = {
        "code_action": "correct-text",
        "correct_value": correct_value,
    }
    if proposed_title:
        r["proposed_title"] = proposed_title
    return r


def propose_insert_text_on_line_after_diagnostic_quick_fix(
    text_to_insert: str,
) -> InsertTextOnLineAfterDiagnosticCodeAction:
    return {
        "code_action": "insert-text-on-line-after-diagnostic",
        "text_to_insert": text_to_insert,
    }


def propose_remove_line_quick_fix() -> RemoveLineCodeAction:
    return {
        "code_action": "remove-line",
    }


def propose_remove_range_quick_fix(
    *,
    proposed_title: str | None = None,
) -> RemoveRangeCodeAction:
    r: RemoveRangeCodeAction = {
        "code_action": "remove-range",
    }
    if proposed_title:
        r["proposed_title"] = proposed_title
    return r


CODE_ACTION_HANDLERS: dict[
    CodeActionName,
    Callable[
        ["LintState", Mapping[str, str], types.CodeActionParams, types.Diagnostic],
        Iterable[types.CodeAction | types.Command],
    ],
] = {}
M = TypeVar("M", bound=Mapping[str, str])
Handler = Callable[
    ["LintState", M, types.CodeActionParams, types.Diagnostic],
    Iterable[Union[types.CodeAction, types.Command]],
]


def _code_handler_for(action_name: CodeActionName) -> Callable[[Handler], Handler]:
    def _wrapper(func: Handler) -> Handler:
        assert action_name not in CODE_ACTION_HANDLERS
        CODE_ACTION_HANDLERS[action_name] = func
        return func

    return _wrapper


@_code_handler_for("correct-text")
def _correct_value_code_action(
    lint_state: "LintState",
    code_action_data: CorrectTextCodeAction,
    code_action_params: types.CodeActionParams,
    diagnostic: types.Diagnostic,
) -> Iterable[types.CodeAction | types.Command]:
    corrected_value = code_action_data["correct_value"]
    title = code_action_data.get("proposed_title")
    if not title:
        title = f'Replace with "{corrected_value}"'
    yield _simple_quick_fix(
        lint_state,
        code_action_params,
        title,
        [diagnostic],
        diagnostic.range,
        corrected_value,
    )


@_code_handler_for("insert-text-on-line-after-diagnostic")
def _insert_text_on_line_after_diagnostic_code_action(
    lint_state: "LintState",
    code_action_data: InsertTextOnLineAfterDiagnosticCodeAction,
    code_action_params: types.CodeActionParams,
    diagnostic: types.Diagnostic,
) -> Iterable[types.CodeAction | types.Command]:
    corrected_value = code_action_data["text_to_insert"]
    line_no = diagnostic.range.end.line
    if diagnostic.range.end.character > 0:
        line_no += 1
    insert_range = types.Range(
        types.Position(
            line_no,
            0,
        ),
        types.Position(
            line_no,
            0,
        ),
    )
    yield _simple_quick_fix(
        lint_state,
        code_action_params,
        f'Insert "{corrected_value}"',
        [diagnostic],
        insert_range,
        corrected_value,
    )


def _simple_quick_fix(
    lint_state: "LintState",
    _code_action_params: types.CodeActionParams,
    title: str,
    diagnostics: list[types.Diagnostic],
    affected_range: types.Range,
    replacement_text: str,
) -> types.CodeAction:
    doc_uri = lint_state.doc_uri
    edit = types.TextEdit(
        affected_range,
        replacement_text,
    )
    if lint_state.workspace_text_edit_support.supports_document_changes:
        ws_edit = types.WorkspaceEdit(
            document_changes=[
                types.TextDocumentEdit(
                    text_document=types.OptionalVersionedTextDocumentIdentifier(
                        doc_uri,
                        lint_state.doc_version,
                    ),
                    edits=[edit],
                )
            ]
        )
    else:
        ws_edit = types.WorkspaceEdit(
            changes={doc_uri: [edit]},
        )
    return types.CodeAction(
        title=title,
        kind=types.CodeActionKind.QuickFix,
        diagnostics=diagnostics,
        edit=ws_edit,
    )


def range_compatible_with_remove_line_fix(range_: types.Range | TERange) -> bool:
    if isinstance(range_, TERange):
        start = range_.start_pos
        end = range_.end_pos
        if start.line_position != end.line_position and (
            start.line_position + 1 != end.line_position or end.cursor_position > 0
        ):
            return False
    else:
        start = range_.start
        end = range_.end
        if start.line != end.line and (start.line + 1 != end.line or end.character > 0):
            return False
    return True


@_code_handler_for("remove-line")
def _remove_line_code_action(
    lint_state: "LintState",
    _code_action_data: RemoveLineCodeAction,
    code_action_params: types.CodeActionParams,
    diagnostic: types.Diagnostic,
) -> Iterable[types.CodeAction | types.Command]:
    start = code_action_params.range.start
    if range_compatible_with_remove_line_fix(code_action_params.range):
        _warn(
            "Bug: the quick was used for a diagnostic that spanned multiple lines and would corrupt the file."
        )
        return

    delete_range = types.Range(
        start=types.Position(
            line=start.line,
            character=0,
        ),
        end=types.Position(
            line=start.line + 1,
            character=0,
        ),
    )
    yield _simple_quick_fix(
        lint_state,
        code_action_params,
        "Remove the line",
        [diagnostic],
        delete_range,
        "",
    )


@_code_handler_for("remove-range")
def _remove_range_code_action(
    lint_state: "LintState",
    code_action_data: RemoveRangeCodeAction,
    code_action_params: types.CodeActionParams,
    diagnostic: types.Diagnostic,
) -> Iterable[types.CodeAction | types.Command]:
    title = code_action_data.get("proposed_title", "Delete")
    yield _simple_quick_fix(
        lint_state,
        code_action_params,
        title,
        [diagnostic],
        diagnostic.range,
        "",
    )


def accepts_quickfixes(
    code_action_params: types.CodeActionParams,
) -> bool:
    only = code_action_params.context.only
    if not only:
        return True
    return types.CodeActionKind.QuickFix in only


def provide_standard_quickfixes_from_diagnostics_ls(
    ls: "DebputyLanguageServer",
    code_action_params: types.CodeActionParams,
) -> list[types.Command | types.CodeAction] | None:
    if not accepts_quickfixes(code_action_params):
        return None
    doc_uri = code_action_params.text_document.uri
    matched_diagnostics = ls.diagnostics_in_range(
        doc_uri,
        code_action_params.range,
    )
    if not matched_diagnostics:
        return None
    doc = ls.workspace.get_text_document(doc_uri)
    return _provide_standard_quickfixes_from_diagnostics(
        ls.lint_state(doc),
        code_action_params,
        matched_diagnostics,
    )


def provide_standard_quickfixes_from_diagnostics_lint(
    lint_state: "LintState",
    code_action_params: types.CodeActionParams,
) -> list[types.Command | types.CodeAction] | None:
    return _provide_standard_quickfixes_from_diagnostics(
        lint_state,
        code_action_params,
        code_action_params.context.diagnostics,
    )


def _provide_standard_quickfixes_from_diagnostics(
    lint_state: "LintState",
    code_action_params: types.CodeActionParams,
    matched_diagnostics: list[types.Diagnostic],
) -> list[types.Command | types.CodeAction] | None:
    actions: list[types.Command | types.CodeAction] = []
    for diagnostic in matched_diagnostics:
        if not isinstance(diagnostic.data, dict):
            continue
        data: DiagnosticData = cast("DiagnosticData", diagnostic.data)
        quickfixes = data.get("quickfixes")
        if quickfixes is None:
            continue
        for action_suggestion in quickfixes:
            if (
                action_suggestion
                and isinstance(action_suggestion, Mapping)
                and "code_action" in action_suggestion
            ):
                action_name: CodeActionName = action_suggestion["code_action"]
                handler = CODE_ACTION_HANDLERS.get(action_name)
                if handler is not None:
                    actions.extend(
                        handler(
                            lint_state,
                            cast("Mapping[str, str]", action_suggestion),
                            code_action_params,
                            diagnostic,
                        )
                    )
                else:
                    _warn(f"No codeAction handler for {action_name} !?")
    if not actions:
        return None
    return actions