import asyncio
import json
import logging
import os
import pathlib
import sys
import traceback
import typing
import warnings
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Type
from typing import Union

from lsprotocol import types
from lsprotocol.converters import get_converter
from packaging.version import parse as parse_version
from pygls.exceptions import JsonRpcException
from pygls.exceptions import PyglsError
from pygls.lsp.client import BaseLanguageClient
from pygls.protocol import default_converter

from .checks import LspSpecificationWarning
from .protocol import LanguageClientProtocol

if sys.version_info.minor < 9:
    import importlib_resources as resources
else:
    import importlib.resources as resources  # type: ignore[no-redef]


__version__ = "0.4.0"
logger = logging.getLogger(__name__)


class LanguageClient(BaseLanguageClient):
    """Used to drive language servers under test."""

    protocol: LanguageClientProtocol

    def __init__(self, *args, configuration: Optional[Dict[str, Any]] = None, **kwargs):
        if "protocol_cls" not in kwargs:
            kwargs["protocol_cls"] = LanguageClientProtocol

        super().__init__("pytest-lsp-client", __version__, *args, **kwargs)

        self.capabilities: Optional[types.ClientCapabilities] = None
        """The client's capabilities."""

        self.shown_documents: List[types.ShowDocumentParams] = []
        """Holds any received show document requests."""

        self.messages: List[types.ShowMessageParams] = []
        """Holds any received ``window/showMessage`` requests."""

        self.log_messages: List[types.LogMessageParams] = []
        """Holds any received ``window/logMessage`` requests."""

        self.diagnostics: Dict[str, List[types.Diagnostic]] = {}
        """Holds any recieved diagnostics."""

        self.progress_reports: Dict[
            types.ProgressToken, List[types.ProgressParams]
        ] = {}
        """Holds any received progress updates."""

        self.error: Optional[Exception] = None
        """Indicates if the client encountered an error."""

        config = (configuration or {"": {}}).copy()
        if "" not in config:
            config[""] = {}

        self._configuration: Dict[str, Dict[str, Any]] = config
        """Holds ``workspace/configuration`` values."""

        self._setup_log_index = 0
        """Used to keep track of which log messages occurred during startup."""

        self._last_log_index = 0
        """Used to keep track of which log messages correspond with which test case."""

    async def server_exit(self, server: asyncio.subprocess.Process):
        """Called when the server process exits."""
        logger.debug("Server process exited with code: %s", server.returncode)

        if self._stop_event.is_set():
            return

        stderr = ""
        if server.stderr is not None:
            stderr_bytes = await server.stderr.read()
            stderr = stderr_bytes.decode("utf8")

        loop = asyncio.get_running_loop()
        loop.call_soon(
            cancel_all_tasks,
            f"Server process exited with return code: {server.returncode}\n{stderr}",
        )

    def report_server_error(
        self, error: Exception, source: Union[PyglsError, JsonRpcException]
    ):
        """Called when the server does something unexpected, e.g. sending malformed
        JSON."""
        self.error = error
        tb = "".join(traceback.format_exc())

        message = f"{source.__name__}: {error}\n{tb}"  # type: ignore

        loop = asyncio.get_running_loop()
        loop.call_soon(cancel_all_tasks, message)

        if self._stop_event:
            self._stop_event.set()

    def get_configuration(
        self, *, section: Optional[str] = None, scope_uri: Optional[str] = None
    ) -> Optional[Any]:
        """Get a configuration value.

        Parameters
        ----------
        section
           The optional section name to retrieve.
           If ``None`` the top level configuration object for the requested scope will
           be returned

        scope_uri
           The scope at which to set the configuration.
           If ``None``, this will default to the global scope.

        Returns
        -------
        Optional[Any]
           The requested configuration value or ``None`` if not found.
        """
        section = section or ""
        scope = scope_uri or ""

        # Find the longest prefix of ``scope``. The empty string is a prefix of all
        # strings so there will always be at least one match
        candidates = [c for c in self._configuration.keys() if scope.startswith(c)]
        selected = sorted(candidates, key=len, reverse=True)[0]

        if (item := self._configuration.get(selected, None)) is None:
            return None

        if section == "":
            return item

        for segment in section.split("."):
            if not hasattr(item, "get"):
                return None

            if (item := item.get(segment, None)) is None:
                return None

        return item

    def set_configuration(
        self,
        item: Any,
        *,
        section: Optional[str] = None,
        scope_uri: Optional[str] = None,
    ):
        """Set a configuration value.

        Parameters
        ----------
        item
           The value to set

        section
           The optional section name to set.
           If ``None`` the top level configuration object will be overriden with
           ``item``.

        scope_uri
           The scope at which to set the configuration.
           If ``None``, this will default to the global scope.
        """
        section = section or ""
        scope = scope_uri or ""

        if section == "":
            self._configuration[scope] = item
            return

        config = self._configuration.setdefault(scope, {})
        *parents, name = section.split(".")

        for segment in parents:
            config = config.setdefault(segment, {})

        config[name] = item

    async def initialize_session(
        self, params: types.InitializeParams
    ) -> types.InitializeResult:
        """Make an ``initialize`` request to a lanaguage server.

        It will also automatically send an ``initialized`` notification once
        the server responds.

        Parameters
        ----------
        params
           The parameters to send to the client.

           The following fields will be automatically set if left blank.

           - ``process_id``: Set to the PID of the current process.

        Returns
        -------
        InitializeResult
           The result received from the client.
        """
        self.capabilities = params.capabilities

        if params.process_id is None:
            params.process_id = os.getpid()

        response = await self.initialize_async(params)
        self.initialized(types.InitializedParams())

        return response

    async def shutdown_session(self) -> None:
        """Shutdown the server under test.

        Helper method that handles sending ``shutdown`` and ``exit`` messages in the
        correct order.

        .. note::

           This method will not attempt to send these messages if a fatal error has
           occurred.

        """
        if self.error is not None or self.capabilities is None:
            return

        await self.shutdown_async(None)
        self.exit(None)

    async def wait_for_notification(self, method: str):
        """Block until a notification with the given method is received.

        Parameters
        ----------
        method
           The notification method to wait for, e.g. ``textDocument/publishDiagnostics``
        """
        return await self.protocol.wait_for_notification_async(method)


def cancel_all_tasks(message: str):
    """Called to cancel all awaited tasks."""

    for task in asyncio.all_tasks():
        if sys.version_info.minor < 9:
            task.cancel()
        else:
            task.cancel(message)


def make_test_lsp_client() -> LanguageClient:
    """Construct a new test client instance with the handlers needed to capture
    additional responses from the server."""

    client = LanguageClient(
        converter_factory=default_converter,
    )

    @client.feature(types.WORKSPACE_CONFIGURATION)
    def configuration(client: LanguageClient, params: types.ConfigurationParams):
        return [
            client.get_configuration(section=item.section, scope_uri=item.scope_uri)
            for item in params.items
        ]

    @client.feature(types.TEXT_DOCUMENT_PUBLISH_DIAGNOSTICS)
    def publish_diagnostics(
        client: LanguageClient, params: types.PublishDiagnosticsParams
    ):
        client.diagnostics[params.uri] = params.diagnostics

    @client.feature(types.WINDOW_WORK_DONE_PROGRESS_CREATE)
    def create_work_done_progress(
        client: LanguageClient, params: types.WorkDoneProgressCreateParams
    ):
        if params.token in client.progress_reports:
            # TODO: Send an error reponse to the client - might require changes
            #       to pygls...
            warnings.warn(
                f"Duplicate progress token: {params.token!r}", LspSpecificationWarning
            )

        client.progress_reports.setdefault(params.token, [])
        return None

    @client.feature(types.PROGRESS)
    def progress(client: LanguageClient, params: types.ProgressParams):
        if params.token not in client.progress_reports:
            warnings.warn(
                f"Unknown progress token: {params.token!r}", LspSpecificationWarning
            )

        if not params.value:
            return

        if (kind := params.value.get("kind", None)) == "begin":
            type_: Type[Any] = types.WorkDoneProgressBegin
        elif kind == "report":
            type_ = types.WorkDoneProgressReport
        elif kind == "end":
            type_ = types.WorkDoneProgressEnd
        else:
            raise TypeError(f"Unknown progress kind: {kind!r}")

        value = client.protocol._converter.structure(params.value, type_)
        client.progress_reports.setdefault(params.token, []).append(value)

    @client.feature(types.WINDOW_LOG_MESSAGE)
    def log_message(client: LanguageClient, params: types.LogMessageParams):
        client.log_messages.append(params)

        levels = [logger.error, logger.warning, logger.info, logger.debug]
        levels[params.type.value - 1](params.message)

    @client.feature(types.WINDOW_SHOW_MESSAGE)
    def show_message(client: LanguageClient, params):
        client.messages.append(params)

    @client.feature(types.WINDOW_SHOW_DOCUMENT)
    def show_document(
        client: LanguageClient, params: types.ShowDocumentParams
    ) -> types.ShowDocumentResult:
        client.shown_documents.append(params)
        return types.ShowDocumentResult(success=True)

    return client


def client_capabilities(client_spec: str) -> types.ClientCapabilities:
    """Find the capabilities that correspond to the given client spec.

    This function supports the following syntax

    ``client-name`` or ``client-name@latest``
       Return the capabilities of the latest version of ``client-name``

    ``client-name@v2``
       Return the latest release of the ``v2`` of ``client-name``

    ``client-name@v2.3.1``
       Return exactly ``v2.3.1`` of ``client-name``

    Parameters
    ----------
    client_spec
       The string describing the client to load the corresponding
       capabilities for.

    Raises
    ------
    ValueError
       If the requested client's capabilities could not be found

    Returns
    -------
    ClientCapabilities
       The requested client capabilities
    """

    candidates: Dict[str, pathlib.Path] = {}

    client_spec = client_spec.replace("-", "_")
    target_version = "latest"

    if "@" in client_spec:
        client_spec, target_version = client_spec.split("@")
        if target_version.startswith("v"):
            target_version = target_version[1:]

    for resource in resources.files("pytest_lsp.clients").iterdir():
        filename = typing.cast(pathlib.Path, resource)

        # Skip the README or any other files that we don't care about.
        if not filename.suffix == ".json":
            continue

        name, version = filename.stem.split("_v")
        if name == client_spec:
            if version.startswith(target_version) or target_version == "latest":
                candidates[version] = filename

    if len(candidates) == 0:
        raise ValueError(
            f"Could not find capabilities for '{client_spec}@{target_version}'"
        )

    # Out of the available candidates, choose the latest version
    selected_version = sorted(candidates.keys(), key=parse_version, reverse=True)[0]
    filename = candidates[selected_version]

    converter = get_converter()
    capabilities = json.loads(filename.read_text())

    params = converter.structure(capabilities, types.InitializeParams)
    logger.info(
        "Selected %s v%s",
        params.client_info.name,  # type: ignore[union-attr]
        params.client_info.version,  # type: ignore[union-attr]
    )

    return params.capabilities
