# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import functools
import json
from pathlib import Path
from typing import Callable, Mapping, Optional, Type, TypeVar

import testslide

from ... import json_rpc
from ...tests import setup
from ..connections import (
    AsyncBytesWriter,
    AsyncTextWriter,
    create_memory_text_reader,
    MemoryBytesWriter,
)
from ..protocol import (
    ClientCapabilities,
    ContentChange,
    DiagnosticTag,
    DidChangeTextDocumentParameters,
    DidCloseTextDocumentParameters,
    DidOpenTextDocumentParameters,
    DidSaveTextDocumentParameters,
    DocumentUri,
    Info,
    InitializationOptions,
    InitializeParameters,
    PublishDiagnosticsClientCapabilities,
    PublishDiagnosticsClientTagSupport,
    PyreHoverResponse,
    read_json_rpc,
    ReadChannelClosedError,
    ShowStatusRequestClientCapabilities,
    TextDocumentClientCapabilities,
    TextDocumentIdentifier,
    TextDocumentItem,
    TextDocumentSyncClientCapabilities,
    WindowClientCapabilities,
    write_json_rpc,
    write_json_rpc_ignore_connection_error,
)

T = TypeVar("T")


class ExceptionRaisingBytesWriter(AsyncBytesWriter):
    """
    An AsyncBytesWriter that always raises a given except when write is invoked.
    """

    def __init__(self, exception: Exception) -> None:
        self.exception = exception

    async def write(self, data: bytes) -> None:
        raise self.exception

    async def close(self) -> None:
        pass


class DocumentUriTest(testslide.TestCase):
    def test_to_file_path(self) -> None:
        def assert_file_path(uri: str, expected: Optional[Path]) -> None:
            self.assertEqual(DocumentUri.parse(uri).to_file_path(), expected)

        assert_file_path("file:///foo/bar", expected=Path("/foo/bar"))
        assert_file_path("file:///foo/space%20%3Fbar", expected=Path("/foo/space ?bar"))
        assert_file_path("file:///foo/bar#frag", expected=Path("/foo/bar"))
        assert_file_path("file://localhost/etc/fstab", expected=Path("/etc/fstab"))
        assert_file_path(
            "file://hostname/path/to/the%20file.txt",
            expected=Path("/path/to/the file.txt"),
        )
        assert_file_path("https://pyre-check.org/", expected=None)
        assert_file_path("nfs://server//a/b/c", expected=None)
        assert_file_path("ssh://foo@192.168.1.1", expected=None)

    def test_from_file_path(self) -> None:
        def assert_file_path(path: Path, expected: str) -> None:
            self.assertEqual(DocumentUri.from_file_path(path).unparse(), expected)

        assert_file_path(Path("/foo/bar"), expected="file:///foo/bar")
        assert_file_path(Path("/foo/space ?bar"), expected="file:///foo/space%20%3Fbar")


class LSPInputOutputTest(testslide.TestCase):
    @setup.async_test
    async def test_read_json_rpc(self) -> None:
        async def assert_parses(input: str, expected: json_rpc.Request) -> None:
            actual = await read_json_rpc(create_memory_text_reader(input))
            self.assertEqual(actual, expected)

        async def assert_not_parsed(
            input: str, exception_type: Type[Exception] = json_rpc.ParseError
        ) -> None:
            with self.assertRaises(exception_type):
                await read_json_rpc(create_memory_text_reader(input))

        await assert_not_parsed("", exception_type=ReadChannelClosedError)
        await assert_not_parsed("derp")
        await assert_not_parsed("Invalid-Header: \r\n\r\n{}")
        await assert_not_parsed("Not-Content-Length: 42\r\n\r\n{}")
        await assert_not_parsed("Content-Length: derp\r\n\r\n{}")
        await assert_not_parsed(
            'Content-Length: 4\r\n\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}'
        )
        await assert_parses(
            'Content-Length: 27\r\n\r\n{"jsonrpc": "2.0", "id": 0}'
            'Content-Length: 44\r\n\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}',
            expected=json_rpc.Request(id=0, method="foo"),
        )
        await assert_parses(
            'Content-Length: 44\r\n\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}',
            expected=json_rpc.Request(id=0, method="foo"),
        )
        await assert_parses(
            'CONTENT-LENGTH: 44\r\n\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}',
            expected=json_rpc.Request(id=0, method="foo"),
        )
        await assert_parses(
            (
                "Content-Length: 44\r\n"
                "Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n"
                '\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}'
            ),
            expected=json_rpc.Request(id=0, method="foo"),
        )
        await assert_parses(
            (
                "Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n"
                "Content-Length: 44\r\n"
                '\r\n{"jsonrpc": "2.0", "id": 0, "method": "foo"}'
            ),
            expected=json_rpc.Request(id=0, method="foo"),
        )

    @setup.async_test
    async def test_write_lsp(self) -> None:
        async def assert_write(response: json_rpc.Response, expected: str) -> None:
            bytes_writer = MemoryBytesWriter()
            await write_json_rpc(AsyncTextWriter(bytes_writer), response)
            actual = bytes_writer.items()[0].decode("utf-8")
            self.assertEqual(actual, expected)

        await assert_write(
            json_rpc.SuccessResponse(id=0, result=42),
            expected=(
                "Content-Length: 41\r\n\r\n"
                + json.dumps({"jsonrpc": "2.0", "id": 0, "result": 42})
            ),
        )
        await assert_write(
            json_rpc.ErrorResponse(id=0, code=42, message="derp"),
            expected=(
                "Content-Length: 69\r\n\r\n"
                + json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": 0,
                        "error": {"code": 42, "message": "derp"},
                    }
                )
            ),
        )

    @setup.async_test
    async def test_write_json_rpc_ignore_connection_error(self) -> None:
        # This invocation should not raise
        await write_json_rpc_ignore_connection_error(
            AsyncTextWriter(ExceptionRaisingBytesWriter(ConnectionResetError())),
            json_rpc.ErrorResponse(
                id=None,
                code=42,
                message="dummy message",
            ),
        )


class LSPHoverResponseTest(testslide.TestCase):
    def test_to_lsp_response(self) -> None:
        result = PyreHoverResponse("type", "docstring").to_lsp_hover_response()
        expected = "```\ntype\n```\ndocstring"
        self.assertEqual(result.contents, expected)

    def test_to_lsp_response_empty(self) -> None:
        result = PyreHoverResponse().to_lsp_hover_response()
        expected = ""
        self.assertEqual(result.contents, expected)


class LSPParsingTest(testslide.TestCase):
    def assert_parsed(
        self,
        parser: Callable[[json_rpc.Parameters], T],
        parameters: Mapping[str, object],
        expected: T,
    ) -> None:
        self.assertEqual(
            parser(json_rpc.ByNameParameters(parameters)),
            expected,
        )

    def assert_not_parsed(
        self,
        parser: Callable[[json_rpc.Parameters], T],
        parameters: Mapping[str, object],
    ) -> None:
        with self.assertRaises(json_rpc.InvalidRequestError):
            parser(json_rpc.ByNameParameters(parameters))

    def test_parse_initialize(self) -> None:
        assert_parsed = functools.partial(
            self.assert_parsed, InitializeParameters.from_json_rpc_parameters
        )
        assert_not_parsed = functools.partial(
            self.assert_not_parsed, InitializeParameters.from_json_rpc_parameters
        )

        assert_not_parsed({})
        assert_not_parsed({"no_capabilities": 42})
        assert_not_parsed({"capabilities": True})
        assert_not_parsed({"processId": "derp"})
        assert_not_parsed({"capabilities": {"textDocument": "foo"}})
        assert_not_parsed(
            {
                "capabilities": {
                    "textDocument": {
                        "publishDiagnostics": {
                            "tagSupport": {"valueSet": [42]},
                        }
                    }
                }
            }
        )

        assert_parsed(
            {
                "processId": 42,
                "clientInfo": {"name": "foo", "version": "v0"},
                "rootUri": "file:///not_relevant",
                "initializationOptions": {"notebookNumber": 12345},
                "capabilities": {
                    "workspace": {},
                    "textDocument": {
                        "publishDiagnostics": {
                            "relatedInformation": True,
                            "versionSupport": False,
                            "tagSupport": {"valueSet": [1, 2]},
                        },
                        "synchronization": {
                            "dynamicRegistration": True,
                            "willSave": True,
                            "willSaveWaitUntil": True,
                            "didSave": True,
                        },
                        "completion": {},
                        "hover": {},
                    },
                    "window": {
                        "workDoneProgress": True,
                        "status": {"dynamicRegistration": False},
                    },
                },
                "trace": "off",
                "workspaceFolders": [
                    {"uri": "file:///workspalce_folder", "name": "foo"}
                ],
            },
            InitializeParameters(
                process_id=42,
                client_info=Info(name="foo", version="v0"),
                capabilities=ClientCapabilities(
                    text_document=TextDocumentClientCapabilities(
                        synchronization=TextDocumentSyncClientCapabilities(
                            did_save=True
                        ),
                        publish_diagnostics=PublishDiagnosticsClientCapabilities(
                            related_information=True,
                            tag_support=PublishDiagnosticsClientTagSupport(
                                value_set=[
                                    DiagnosticTag.UNNECESSARY,
                                    DiagnosticTag.DEPRECATED,
                                ]
                            ),
                            version_support=False,
                        ),
                    ),
                    window=WindowClientCapabilities(
                        work_done_progress=True,
                        status=ShowStatusRequestClientCapabilities(),
                    ),
                ),
                initialization_options=InitializationOptions(notebook_number=12345),
            ),
        )

    def test_parse_did_open(self) -> None:
        assert_parsed = functools.partial(
            self.assert_parsed, DidOpenTextDocumentParameters.from_json_rpc_parameters
        )
        assert_not_parsed = functools.partial(
            self.assert_not_parsed,
            DidOpenTextDocumentParameters.from_json_rpc_parameters,
        )

        assert_not_parsed({})
        assert_not_parsed({"no_text_document": 42})
        assert_not_parsed({"textDocument": "derp"})
        assert_not_parsed(
            {"textDocument": {"missing_uri": "foo", "version": 0, "text": ""}}
        )
        assert_parsed(
            {
                "textDocument": {
                    "languageId": "python",
                    "text": "foo",
                    "uri": "file:///home/user/test.py",
                    "version": 0,
                }
            },
            expected=DidOpenTextDocumentParameters(
                text_document=TextDocumentItem(
                    language_id="python",
                    text="foo",
                    uri="file:///home/user/test.py",
                    version=0,
                )
            ),
        )

    def test_parse_did_close(self) -> None:
        assert_parsed = functools.partial(
            self.assert_parsed, DidCloseTextDocumentParameters.from_json_rpc_parameters
        )
        assert_not_parsed = functools.partial(
            self.assert_not_parsed,
            DidCloseTextDocumentParameters.from_json_rpc_parameters,
        )

        assert_not_parsed({})
        assert_not_parsed({"no_text_document": 42})
        assert_not_parsed({"textDocument": "derp"})
        assert_not_parsed({"textDocument": {"missing_uri": "foo"}})
        assert_parsed(
            {
                "textDocument": {
                    "uri": "file:///home/user/test.py",
                }
            },
            expected=DidCloseTextDocumentParameters(
                text_document=TextDocumentIdentifier(
                    uri="file:///home/user/test.py",
                )
            ),
        )

    def test_parse_did_save(self) -> None:
        assert_parsed = functools.partial(
            self.assert_parsed, DidSaveTextDocumentParameters.from_json_rpc_parameters
        )
        assert_not_parsed = functools.partial(
            self.assert_not_parsed,
            DidSaveTextDocumentParameters.from_json_rpc_parameters,
        )

        assert_not_parsed({})
        assert_not_parsed({"no_text_document": 42})

        assert_parsed(
            {
                "textDocument": {
                    "uri": "file:///home/user/test.py",
                }
            },
            expected=DidSaveTextDocumentParameters(
                text_document=TextDocumentIdentifier(uri="file:///home/user/test.py")
            ),
        )
        assert_parsed(
            {
                "textDocument": {
                    "uri": "file:///home/user/test.py",
                },
                "text": "foo",
            },
            expected=DidSaveTextDocumentParameters(
                text_document=TextDocumentIdentifier(uri="file:///home/user/test.py"),
                text="foo",
            ),
        )

    def test_parse_did_change(self) -> None:
        assert_parsed = functools.partial(
            self.assert_parsed, DidChangeTextDocumentParameters.from_json_rpc_parameters
        )
        assert_not_parsed = functools.partial(
            self.assert_not_parsed,
            DidChangeTextDocumentParameters.from_json_rpc_parameters,
        )

        assert_not_parsed({})
        assert_not_parsed({"no_text_document": 42})
        assert_not_parsed({"textDocument": "derp"})

        assert_parsed(
            {
                "textDocument": {
                    "uri": "file:///home/user/test.py",
                },
                "contentChanges": [{"text": "x = 5"}],
            },
            expected=DidChangeTextDocumentParameters(
                text_document=TextDocumentIdentifier(uri="file:///home/user/test.py"),
                content_changes=[ContentChange("x = 5")],
            ),
        )
        assert_parsed(
            {
                "textDocument": {
                    "uri": "file:///home/user/test.py",
                },
                "contentChanges": [{"text": "x = 5"}, {"text": "file2: str"}],
            },
            expected=DidChangeTextDocumentParameters(
                text_document=TextDocumentIdentifier(uri="file:///home/user/test.py"),
                content_changes=[ContentChange("x = 5"), ContentChange("file2: str")],
            ),
        )
