from __future__ import annotations

import sys
import time
from multiprocessing import Queue, get_context
from unittest import TestCase, main

import pytest

from mypy.ipc import IPCClient, IPCServer

CONNECTION_NAME = "dmypy-test-ipc"


def server(msg: str, q: Queue[str]) -> None:
    server = IPCServer(CONNECTION_NAME)
    q.put(server.connection_name)
    data = ""
    while not data:
        with server:
            server.write(msg)
            data = server.read()
    server.cleanup()


def server_multi_message_echo(q: Queue[str]) -> None:
    server = IPCServer(CONNECTION_NAME)
    q.put(server.connection_name)
    data = ""
    with server:
        while data != "quit":
            data = server.read()
            server.write(data)
    server.cleanup()


def server_with_short_timeout(msg: str, q: Queue[str]) -> None:
    """Serve a client that takes longer to speak than the connection timeout.

    The timeout bounds accepting a connection, not the traffic after it, so the
    read below must wait however long the client needs.
    """
    server = IPCServer(CONNECTION_NAME, timeout=1)
    q.put(server.connection_name)
    with server:
        data = server.read()
        server.write(data + msg)
    server.cleanup()


class IPCTests(TestCase):
    def setUp(self) -> None:
        if sys.platform == "linux":
            # The default "fork" start method is potentially unsafe
            self.ctx = get_context("forkserver")
        else:
            self.ctx = get_context("spawn")

    def test_transaction_large(self) -> None:
        queue: Queue[str] = self.ctx.Queue()
        msg = "t" * 200000  # longer than the max read size of 100_000
        p = self.ctx.Process(target=server, args=(msg, queue), daemon=True)
        p.start()
        connection_name = queue.get()
        with IPCClient(connection_name, timeout=1) as client:
            assert client.read() == msg
            client.write("test")
        queue.close()
        queue.join_thread()
        p.join()

    def test_connect_twice(self) -> None:
        queue: Queue[str] = self.ctx.Queue()
        msg = "this is a test message"
        p = self.ctx.Process(target=server, args=(msg, queue), daemon=True)
        p.start()
        connection_name = queue.get()
        with IPCClient(connection_name, timeout=1) as client:
            assert client.read() == msg
            client.write("")  # don't let the server hang up yet, we want to connect again.

        with IPCClient(connection_name, timeout=1) as client:
            assert client.read() == msg
            client.write("test")
        queue.close()
        queue.join_thread()
        p.join()
        assert p.exitcode == 0

    def test_multiple_messages(self) -> None:
        queue: Queue[str] = self.ctx.Queue()
        p = self.ctx.Process(target=server_multi_message_echo, args=(queue,), daemon=True)
        p.start()
        connection_name = queue.get()
        with IPCClient(connection_name, timeout=1) as client:
            # "foo bar" with extra accents on letters.
            # In UTF-8 encoding so we don't confuse editors opening this file.
            fancy_text = b"f\xcc\xb6o\xcc\xb2\xf0\x9d\x91\x9c \xd0\xb2\xe2\xb7\xa1a\xcc\xb6r\xcc\x93\xcd\x98\xcd\x8c"
            client.write(fancy_text.decode("utf-8"))
            assert client.read() == fancy_text.decode("utf-8")

            client.write("Test with spaces")
            client.write("Test write before reading previous")
            time.sleep(0)  # yield to the server to force reading of all messages by server.
            assert client.read() == "Test with spaces"
            assert client.read() == "Test write before reading previous"

            client.write("quit")
            assert client.read() == "quit"
        queue.close()
        queue.join_thread()
        p.join()
        assert p.exitcode == 0

    def test_server_timeout_does_not_apply_after_connecting(self) -> None:
        # A server's timeout bounds accepting a connection only. A client that is
        # merely slow to speak must not be mistaken for one that has hung up: that
        # is what killed build workers waiting on the coordinator to load the graph
        # (see #21484). POSIX gets this for free because accept() hands back a
        # blocking socket, so this mainly guards the Windows path.
        queue: Queue[str] = self.ctx.Queue()
        msg = " -- echoed"
        p = self.ctx.Process(target=server_with_short_timeout, args=(msg, queue), daemon=True)
        p.start()
        connection_name = queue.get()
        with IPCClient(connection_name, timeout=1) as client:
            # Stay quiet for well past the server's 1s connection timeout.
            time.sleep(2.5)
            client.write("hello")
            assert client.read() == "hello" + msg
        queue.close()
        queue.join_thread()
        p.join()
        assert p.exitcode == 0

    # Run test_connect_twice a lot, in the hopes of finding issues.
    # This is really slow, so it is skipped, but can be enabled if
    # needed to debug IPC issues.
    @pytest.mark.skip
    def test_connect_alot(self) -> None:
        t0 = time.time()
        for i in range(1000):
            try:
                print(i, "start")
                self.test_connect_twice()
            finally:
                t1 = time.time()
                print(i, t1 - t0)
                sys.stdout.flush()
                t0 = t1


if __name__ == "__main__":
    main()