from __future__ import annotations

import asyncio
import gc
import os
import signal
import sys
import threading
import weakref
from time import sleep

import psutil
import pytest
from tornado.ioloop import IOLoop

from distributed.compatibility import LINUX, MACOS, WINDOWS
from distributed.metrics import time
from distributed.process import AsyncProcess
from distributed.utils import get_mp_context
from distributed.utils_test import gen_test, nodebug


def feed(in_q, out_q):
    obj = in_q.get(timeout=5)
    out_q.put(obj)


def exit(q):
    sys.exit(q.get())


def exit_now(rc=0):
    sys.exit(rc)


def exit_with_sigint():
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    os.kill(os.getpid(), signal.SIGINT)
    while True:
        sleep(0.01)


def wait():
    while True:
        sleep(0.01)


def threads_info(q):
    q.put(len(threading.enumerate()))
    q.put(threading.current_thread().name)


@nodebug
@gen_test()
async def test_simple():
    to_child = get_mp_context().Queue()
    from_child = get_mp_context().Queue()

    proc = AsyncProcess(target=feed, args=(to_child, from_child))
    assert not proc.is_alive()
    assert proc.pid is None
    assert proc.exitcode is None
    assert not proc.daemon
    proc.daemon = True
    assert proc.daemon

    wr1 = weakref.ref(proc)
    wr2 = weakref.ref(proc._process)

    # join() before start()
    with pytest.raises(AssertionError):
        await proc.join()

    await proc.start()
    assert proc.is_alive()
    assert proc.pid is not None
    assert proc.exitcode is None

    t1 = time()
    with pytest.raises(asyncio.TimeoutError):
        await proc.join(timeout=0.02)
    dt = time() - t1
    assert 0.2 >= dt >= 0.001
    assert proc.is_alive()
    assert proc.pid is not None
    assert proc.exitcode is None

    # setting daemon attribute after start()
    with pytest.raises(AssertionError):
        proc.daemon = False

    to_child.put(5)
    assert from_child.get() == 5

    # child should be stopping now
    t1 = time()
    await proc.join(timeout=30)
    dt = time() - t1
    assert dt <= 1.0
    assert not proc.is_alive()
    assert proc.pid is not None
    assert proc.exitcode == 0

    # join() again
    t1 = time()
    await proc.join()
    dt = time() - t1
    assert dt <= 0.6

    del proc

    start = time()
    while wr1() is not None and time() < start + 1:
        # Perhaps the GIL switched before _watch_process() exit,
        # help it a little
        sleep(0.001)
        gc.collect()

    if wr1() is not None:
        # Help diagnosing
        from types import FrameType

        p = wr1()
        if p is not None:
            rc = sys.getrefcount(p)
            refs = gc.get_referrers(p)
            del p
            print("refs to proc:", rc, refs)
            frames = [r for r in refs if isinstance(r, FrameType)]
            for i, f in enumerate(frames):
                print(
                    "frames #%d:" % i,
                    f.f_code.co_name,
                    f.f_code.co_filename,
                    sorted(f.f_locals),
                )
        pytest.fail("AsyncProcess should have been destroyed")
    t1 = time()
    while wr2() is not None:
        await asyncio.sleep(0.01)
        gc.collect()
        dt = time() - t1
        assert dt < 2.0


@gen_test()
async def test_exitcode():
    q = get_mp_context().Queue()

    proc = AsyncProcess(target=exit, kwargs={"q": q})
    proc.daemon = True
    assert not proc.is_alive()
    assert proc.exitcode is None

    await proc.start()
    assert proc.is_alive()
    assert proc.exitcode is None

    q.put(5)
    await proc.join()
    assert not proc.is_alive()
    assert proc.exitcode == 5


def assert_exit_code(proc: AsyncProcess, expect: signal.Signals) -> None:
    # Note: WINDOWS constant as doesn't work with `mypy --platform win32`
    if sys.platform == "win32":
        # multiprocessing.Process.terminate() sets exit code -15 like in Linux, but
        # os.kill(pid, signal.SIGTERM) sets exit code +15
        assert proc.exitcode in (-expect, expect)
    elif MACOS:
        # FIXME this happens very frequently on GitHub MacOSX CI. Reason unknown.
        if expect != signal.SIGKILL and proc.exitcode == -signal.SIGKILL:
            raise pytest.xfail(reason="https://github.com/dask/distributed/issues/6393")
        assert proc.exitcode == -expect
    else:
        assert LINUX
        assert proc.exitcode == -expect


@gen_test()
async def test_sigint_from_same_process():
    proc = AsyncProcess(target=exit_with_sigint)
    assert not proc.is_alive()
    assert proc.exitcode is None

    await proc.start()
    await proc.join()

    assert not proc.is_alive()
    assert_exit_code(proc, signal.SIGINT)


@gen_test()
async def test_sigterm_from_parent_process():
    proc = AsyncProcess(target=wait)
    await proc.start()
    os.kill(proc.pid, signal.SIGTERM)
    await proc.join()
    assert not proc.is_alive()
    assert_exit_code(proc, signal.SIGTERM)


@gen_test()
async def test_terminate():
    proc = AsyncProcess(target=wait)
    await proc.start()
    await proc.terminate()
    await proc.join()
    assert not proc.is_alive()
    assert_exit_code(proc, signal.SIGTERM)


@gen_test()
async def test_close():
    proc = AsyncProcess(target=exit_now)
    proc.close()
    with pytest.raises(ValueError):
        await proc.start()

    proc = AsyncProcess(target=exit_now)
    await proc.start()
    proc.close()
    with pytest.raises(ValueError):
        await proc.terminate()

    proc = AsyncProcess(target=exit_now)
    await proc.start()
    await proc.join()
    proc.close()
    with pytest.raises(ValueError):
        await proc.join()
    proc.close()


@gen_test()
async def test_exit_callback():
    to_child = get_mp_context().Queue()
    from_child = get_mp_context().Queue()
    evt = asyncio.Event()

    def on_stop(_proc):
        evt.set()

    # Normal process exit
    proc = AsyncProcess(target=feed, args=(to_child, from_child))
    evt.clear()
    proc.set_exit_callback(on_stop)
    proc.daemon = True

    await proc.start()
    await asyncio.sleep(0.05)
    assert proc.is_alive()
    assert not evt.is_set()

    to_child.put(None)
    await asyncio.wait_for(evt.wait(), 5)
    assert evt.is_set()
    assert not proc.is_alive()

    # Process terminated
    proc = AsyncProcess(target=wait)
    evt.clear()
    proc.set_exit_callback(on_stop)
    proc.daemon = True

    await proc.start()
    await asyncio.sleep(0.05)
    assert proc.is_alive()
    assert not evt.is_set()

    await proc.terminate()
    await asyncio.wait_for(evt.wait(), 5)
    assert evt.is_set()


@gen_test()
async def test_child_main_thread():
    """
    The main thread in the child should be called "MainThread".
    """
    q = get_mp_context().Queue()
    proc = AsyncProcess(target=threads_info, args=(q,))
    await proc.start()
    await proc.join()
    n_threads = q.get()
    main_name = q.get()
    assert n_threads <= 3
    assert main_name == "MainThread"
    q.close()
    q._reader.close()
    q._writer.close()


@pytest.mark.skipif(WINDOWS, reason="num_fds not supported on windows")
@gen_test()
async def test_num_fds():
    # Warm up
    proc = AsyncProcess(target=exit_now)
    proc.daemon = True
    await proc.start()
    await proc.join()

    p = psutil.Process()
    before = p.num_fds()

    proc = AsyncProcess(target=exit_now)
    proc.daemon = True
    await proc.start()
    await proc.join()
    assert not proc.is_alive()
    assert proc.exitcode == 0

    while p.num_fds() > before:
        await asyncio.sleep(0.01)


@gen_test()
async def test_terminate_after_stop():
    proc = AsyncProcess(target=sleep, args=(0,))
    await proc.start()
    await asyncio.sleep(0.1)
    await proc.terminate()
    await proc.join()


def kill_target(ev):
    signal.signal(signal.SIGTERM, signal.SIG_IGN)
    ev.set()
    sleep(300)


@pytest.mark.skipif(WINDOWS, reason="Needs SIGKILL")
@gen_test()
async def test_kill():
    ev = get_mp_context().Event()
    proc = AsyncProcess(target=kill_target, args=(ev,))
    await proc.start()
    ev.wait()
    await proc.kill()
    await proc.join()
    assert not proc.is_alive()
    assert proc.exitcode == -signal.SIGKILL


def _worker_process(worker_ready, child_pipe):
    # child_pipe is the write-side of the children_alive pipe held by the
    # test process. When this _worker_process exits, this file descriptor should
    # have no references remaining anywhere and be closed by the kernel. The
    # test will therefore be able to tell that this process has exited by
    # reading children_alive.

    # Signal to parent process that this process has started and made it this
    # far. This should cause the parent to exit rapidly after this statement.
    worker_ready.set()

    # The parent exiting should cause this process to os._exit from a monitor
    # thread. This sleep should never return.
    shorter_timeout = 2.5  # timeout shorter than that in the spawning test.
    sleep(shorter_timeout)

    # Unreachable if functioning correctly.
    child_pipe.send("child should have exited by now")


def _parent_process(child_pipe):
    """Simulate starting an AsyncProcess and then dying.

    The child_alive pipe is held open for as long as the child is alive, and can
    be used to determine if it exited correctly."""

    async def parent_process_coroutine():
        IOLoop.current()
        worker_ready = get_mp_context().Event()

        worker = AsyncProcess(target=_worker_process, args=(worker_ready, child_pipe))

        await worker.start()

        # Wait for the child process to have started.
        worker_ready.wait()

        # Exit immediately, without doing any process teardown (including atexit
        # and 'finally:' blocks) as if by SIGKILL. This should cause
        # worker_process to also exit.
        os._exit(255)

    async def run_with_timeout():
        t = asyncio.create_task(parent_process_coroutine())
        return await asyncio.wait_for(t, timeout=10)

    asyncio.run(run_with_timeout())
    raise RuntimeError("this should be unreachable due to os._exit")


def test_asyncprocess_child_teardown_on_parent_exit():
    r"""Check that a child process started by AsyncProcess exits if its parent
    exits.

    The motivation is to ensure that if an AsyncProcess is created and the
    creator process dies unexpectedly (e.g, via Out-of-memory SIGKILL), the
    child process and resources held by it should not be leaked.

    The child should monitor its parent and exit promptly if the parent exits.

    [test process] -> [parent using AsyncProcess (dies)] -> [worker process]
                 \                                          /
                  \________ <--   child_pipe   <-- ________/
    """
    # When child_pipe is closed, the children_alive pipe unblocks.
    children_alive, child_pipe = get_mp_context().Pipe(duplex=False)

    try:
        parent = get_mp_context().Process(target=_parent_process, args=(child_pipe,))
        parent.start()

        # Close our reference to child_pipe so that the child has the only one.
        child_pipe.close()

        # Wait for the parent to exit. By the time join returns, the child
        # process is orphaned, and should be in the process of exiting by
        # itself.
        parent.join()

        # By the time we reach here,the parent has exited. The parent only exits
        # when the child is ready to enter the sleep, so all of the slow things
        # (process startup, etc) should have happened by now, even on a busy
        # system. A short timeout should therefore be appropriate.
        short_timeout = 5.0
        # Poll is used to allow other tests to proceed after this one in case of
        # test failure.
        try:
            readable = children_alive.poll(short_timeout)
        except BrokenPipeError:
            assert WINDOWS, "should only raise on windows"
            # Broken pipe implies closed, which is readable.
            readable = True

        # If this assert fires, then something went wrong. Either the child
        # should write into the pipe, or it should exit and the pipe should be
        # closed (which makes it become readable).
        assert readable

        try:
            # This won't block due to the above 'assert readable'.
            result = children_alive.recv()
        except EOFError:
            pass  # Test passes.
        except BrokenPipeError:
            assert WINDOWS, "should only raise on windows"
            # Test passes.
        else:
            # Oops, children_alive read something. It should be closed. If
            # something was read, it's a message from the child telling us they
            # are still alive!
            raise RuntimeError(f"unreachable: {result}")

    finally:
        # Cleanup.
        children_alive.close()