from __future__ import annotations

import os
import threading
import time
from contextlib import contextmanager
from queue import Queue
from typing import TYPE_CHECKING, NoReturn

import pytest

from .. import _thread_cache
from .._thread_cache import ThreadCache, start_thread_soon
from .tutil import gc_collect_harder, slow

if TYPE_CHECKING:
    from collections.abc import Iterator

    from outcome import Outcome


def test_thread_cache_basics() -> None:
    q: Queue[Outcome[object]] = Queue()

    def fn() -> NoReturn:
        raise RuntimeError("hi")

    def deliver(outcome: Outcome[object]) -> None:
        q.put(outcome)

    start_thread_soon(fn, deliver)

    outcome = q.get()
    with pytest.raises(RuntimeError, match="hi"):
        outcome.unwrap()


def test_thread_cache_deref() -> None:
    res = [False]

    class del_me:
        def __call__(self) -> int:
            return 42

        def __del__(self) -> None:
            res[0] = True

    q: Queue[Outcome[int]] = Queue()

    def deliver(outcome: Outcome[int]) -> None:
        q.put(outcome)

    start_thread_soon(del_me(), deliver)
    outcome = q.get()
    assert outcome.unwrap() == 42

    gc_collect_harder()
    assert res[0]


@slow
def test_spawning_new_thread_from_deliver_reuses_starting_thread() -> None:
    # We know that no-one else is using the thread cache, so if we keep
    # submitting new jobs the instant the previous one is finished, we should
    # keep getting the same thread over and over. This tests both that the
    # thread cache is LIFO, and that threads can be assigned new work *before*
    # deliver exits.

    # Make sure there are a few threads running, so if we weren't LIFO then we
    # could grab the wrong one.
    q: Queue[Outcome[object]] = Queue()
    COUNT = 5
    for _ in range(COUNT):
        start_thread_soon(lambda: time.sleep(1), lambda result: q.put(result))
    for _ in range(COUNT):
        q.get().unwrap()

    seen_threads = set()
    done = threading.Event()

    def deliver(n: int, _: object) -> None:
        print(n)
        seen_threads.add(threading.current_thread())
        if n == 0:
            done.set()
        else:
            start_thread_soon(lambda: None, lambda _: deliver(n - 1, _))

    start_thread_soon(lambda: None, lambda _: deliver(5, _))

    done.wait()

    assert len(seen_threads) == 1


@slow
def test_idle_threads_exit(monkeypatch: pytest.MonkeyPatch) -> None:
    # Temporarily set the idle timeout to something tiny, to speed up the
    # test. (But non-zero, so that the worker loop will at least yield the
    # CPU.)
    monkeypatch.setattr(_thread_cache, "IDLE_TIMEOUT", 0.0001)

    q: Queue[threading.Thread] = Queue()
    start_thread_soon(lambda: None, lambda _: q.put(threading.current_thread()))
    seen_thread = q.get()
    # Since the idle timeout is 0, after sleeping for 1 second, the thread
    # should have exited
    time.sleep(1)
    assert not seen_thread.is_alive()


@contextmanager
def _join_started_threads() -> Iterator[None]:
    before = frozenset(threading.enumerate())
    try:
        yield
    finally:
        for thread in threading.enumerate():
            if thread not in before:
                thread.join(timeout=1.0)
                assert not thread.is_alive()


def test_race_between_idle_exit_and_job_assignment(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    # This is a lock where the first few times you try to acquire it with a
    # timeout, it waits until the lock is available and then pretends to time
    # out. Using this in our thread cache implementation causes the following
    # sequence:
    #
    # 1. start_thread_soon grabs the worker thread, assigns it a job, and
    #    releases its lock.
    # 2. The worker thread wakes up (because the lock has been released), but
    #    the JankyLock lies to it and tells it that the lock timed out. So the
    #    worker thread tries to exit.
    # 3. The worker thread checks for the race between exiting and being
    #    assigned a job, and discovers that it *is* in the process of being
    #    assigned a job, so it loops around and tries to acquire the lock
    #    again.
    # 4. Eventually the JankyLock admits that the lock is available, and
    #    everything proceeds as normal.

    class JankyLock:
        def __init__(self) -> None:
            self._lock = threading.Lock()
            self._counter = 3

        def acquire(self, timeout: int = -1) -> bool:
            got_it = self._lock.acquire(timeout=timeout)
            if timeout == -1:
                return True
            elif got_it:
                if self._counter > 0:
                    self._counter -= 1
                    self._lock.release()
                    return False
                return True
            else:
                return False

        def release(self) -> None:
            self._lock.release()

    monkeypatch.setattr(_thread_cache, "Lock", JankyLock)

    with _join_started_threads():
        tc = ThreadCache()
        done = threading.Event()
        tc.start_thread_soon(lambda: None, lambda _: done.set())
        done.wait()
        # Let's kill the thread we started, so it doesn't hang around until the
        # test suite finishes. Doesn't really do any harm, but it can be confusing
        # to see it in debug output.
        monkeypatch.setattr(_thread_cache, "IDLE_TIMEOUT", 0.0001)
        tc.start_thread_soon(lambda: None, lambda _: None)


def test_raise_in_deliver(capfd: pytest.CaptureFixture[str]) -> None:
    seen_threads = set()

    def track_threads() -> None:
        seen_threads.add(threading.current_thread())

    def deliver(_: object) -> NoReturn:
        done.set()
        raise RuntimeError("don't do this")

    done = threading.Event()
    start_thread_soon(track_threads, deliver)
    done.wait()
    done = threading.Event()
    start_thread_soon(track_threads, lambda _: done.set())
    done.wait()
    assert len(seen_threads) == 1
    err = capfd.readouterr().err
    assert "don't do this" in err
    assert "delivering result" in err


@pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork isn't supported")
def test_clear_thread_cache_after_fork() -> None:
    assert hasattr(os, "fork")

    def foo() -> None:
        pass

    # ensure the thread cache exists
    done = threading.Event()
    start_thread_soon(foo, lambda _: done.set())
    done.wait()

    child_pid = os.fork()

    # try using it
    done = threading.Event()
    start_thread_soon(foo, lambda _: done.set())
    done.wait()

    if child_pid != 0:
        # if this test fails, this will hang, triggering a timeout.
        os.waitpid(child_pid, 0)
    else:
        # this is necessary because os._exit doesn't unwind the stack,
        # so coverage doesn't get to automatically stop and save
        # coverage information.
        try:
            import coverage

            cov = coverage.Coverage.current()
            # the following pragmas are necessary because if coverage:
            #  - isn't running, then it can't record the branch not
            #    taken
            #  - isn't installed, then it can't record the ImportError

            if cov:  # pragma: no branch
                cov.stop()
                cov.save()
        except ImportError:  # pragma: no cover
            pass

        os._exit(0)  # pragma: no cover  # coverage was stopped above.
