from concurrent.futures import (  # pylint: disable=redefined-builtin
    ThreadPoolExecutor,
    TimeoutError,
)
from functools import partial
from threading import Event
import sys
import logging
from six.moves.queue import Queue

from pytest import fixture
from hamcrest import (
    assert_that,
    equal_to,
    calling,
    raises,
    has_length,
    has_item,
    contains,
    matches_regexp,
)

from more_executors import PollExecutor

from .util import assert_soon


@fixture
def executor():
    return ThreadPoolExecutor()


if sys.version_info[0:1] < (2, 7):
    # This python is too old for pytest's caplog,
    # make a null caplog and skip that part of the test
    @fixture
    def caplog():
        pass


def poll_tasks(tasks, poll_descriptors):
    for descriptor in poll_descriptors:
        task_id = descriptor.result
        status = tasks.get(task_id)
        if status == "done":
            descriptor.yield_result("done")
        elif status == "error":
            descriptor.yield_exception(RuntimeError("task failed: %s" % task_id))


def test_basic_poll(executor):
    task_id_queue = Queue()
    tasks = {}
    poll_fn = partial(poll_tasks, tasks)
    poll_executor = PollExecutor(executor, poll_fn, default_interval=0.01)

    def make_task(x):
        return "%s-%s" % (x, task_id_queue.get(True))

    inputs = ["a", "b", "c"]
    futures = [poll_executor.submit(make_task, x) for x in inputs]

    # The futures should not currently be able to progress.
    assert_that(not any([f.done() for f in futures]))

    # Allow tasks to be created.
    task_id_queue.put("x")
    task_id_queue.put("y")
    task_id_queue.put("z")

    # Insert some task statuses for the poll function to detect.
    # Note that we can't guess which thread got which queue item,
    # so let's just spam with every combination
    tasks["a-x"] = "done"
    tasks["a-y"] = "done"
    tasks["a-z"] = "done"
    tasks["b-x"] = "error"
    tasks["b-y"] = "error"
    tasks["b-z"] = "error"
    # Leave c with no result

    # Future a should become resolved
    assert_that(futures[0].result(10), equal_to("done"))

    # Future b should raise an exception
    assert_that(
        calling(futures[1].result).with_args(10), raises(RuntimeError, "task failed")
    )

    # Future c should still be waiting
    assert_that(not futures[2].done())


def test_poll_notify(executor):
    task_id_queue = Queue()
    tasks = {}
    poll_fn = partial(poll_tasks, tasks)

    # Make the default interval unreasonably large, so notify() is the only
    # way we'll really poll
    poll_executor = PollExecutor(executor, poll_fn, default_interval=60.0)

    def make_task(x):
        return "%s-%s" % (x, task_id_queue.get(True))

    inputs = ["a", "b", "c"]
    futures = [poll_executor.submit(make_task, x) for x in inputs]

    # Futures should not be able to resolve yet
    assert_that(calling(futures[0].result).with_args(0.1), raises(TimeoutError))
    assert_that(calling(futures[1].result).with_args(0.1), raises(TimeoutError))

    # Allow tasks to be created.
    task_id_queue.put("x")
    task_id_queue.put("y")
    task_id_queue.put("z")

    # Futures should not be able to resolve yet
    assert_that(calling(futures[0].result).with_args(0.1), raises(TimeoutError))
    assert_that(calling(futures[1].result).with_args(0.1), raises(TimeoutError))

    # Insert some task statuses for the poll function to detect.
    # Note that we can't guess which thread got which queue item,
    # so let's just spam with every combination
    tasks["a-x"] = "done"
    tasks["a-y"] = "done"
    tasks["a-z"] = "done"
    tasks["b-x"] = "error"
    tasks["b-y"] = "error"
    tasks["b-z"] = "error"
    # Leave c with no result

    # Futures should still not be able to resolve (between polls)
    assert_that(calling(futures[0].result).with_args(0.1), raises(TimeoutError))
    assert_that(calling(futures[1].result).with_args(0.1), raises(TimeoutError))

    # But if we notify...
    poll_executor.notify()

    # Now they should be able to resolve
    assert_that(futures[0].result(10), equal_to("done"))
    assert_that(
        calling(futures[1].result).with_args(10), raises(RuntimeError, "task failed")
    )

    # Future c should still be waiting
    assert_that(not futures[2].done())


def test_cancel_fn(executor, caplog):
    task_id_queue = Queue()
    tasks = {}
    poll_fn = partial(poll_tasks, tasks)

    def cancel_fn(task):
        if task.startswith("cancel-true-"):
            return True
        if task.startswith("cancel-false-"):
            return False
        raise RuntimeError("simulated cancel error")

    poll_executor = PollExecutor(executor, poll_fn, cancel_fn, default_interval=0.01)

    def make_task(x):
        got = task_id_queue.get(True)
        task_id_queue.task_done()
        return "%s-%s" % (x, got)

    inputs = ["cancel-true", "cancel-false", "cancel-error"]
    futures = [poll_executor.submit(make_task, x) for x in inputs]

    # The futures should not currently be able to progress.
    assert_that(not any([f.done() for f in futures]))

    # Allow tasks to be created.
    task_id_queue.put("x")
    task_id_queue.put("y")
    task_id_queue.put("z")

    # Wait until all tasks were created and futures moved
    # into poll mode
    task_id_queue.join()

    # Wait until the make_task function definitely completed in each thread,
    # which can be determined by running==False
    assert_soon(lambda: assert_that(all([not f.running() for f in futures])))

    # Should be able to cancel soon.
    # Why "soon" instead of "now" - because even though the futures above are
    # not running, the delegate may not have been cleared yet.  Cancel needs
    # to wait until the future's delegate is cleared and the future has
    # transitioned fully into "poll mode".
    assert_soon(lambda: assert_that(futures[0].cancel()))

    # The other two futures don't need assert_soon since the cancel result is negative.

    # Cancel behavior should be consistent (calling multiple times same
    # as calling once)
    for _ in 1, 2:
        # Cancelling the cancel-true task should be allowed.
        assert_that(futures[0].cancel())

        # Cancelling the cancel-false task should not be allowed.
        assert_that(not futures[1].cancel())

        # Cancelling the cancel-error task should not be allowed.
        assert_that(not futures[2].cancel())

    # An error should have been logged due to the cancel function raising.
    if caplog:
        assert_that(
            caplog.record_tuples,
            has_item(
                contains(
                    "PollExecutor",
                    logging.ERROR,
                    matches_regexp(r"Exception during cancel .*/cancel-error"),
                )
            ),
        )


def test_cancel_during_poll(executor):
    task_ran = Event()
    poll_ran = Event()
    last_descriptors = []

    def poll_fn(descriptors):
        while last_descriptors:
            last_descriptors.pop()
        last_descriptors.extend(descriptors)
        poll_ran.set()

    def fn():
        task_ran.set()
        return 123

    poll_executor = PollExecutor(executor, poll_fn, default_interval=0.01)
    future = poll_executor.submit(fn)

    # It shouldn't finish yet.
    assert_that(not future.done())

    # Wait until the delegate has definitely started executing.
    task_ran.wait(10)

    # To determine once the submitted function has completed, we can
    # wait until 'running' is no longer true.
    assert_soon(lambda: assert_that(not future.running()))

    # Wait until next poll
    poll_ran.clear()
    poll_ran.wait(10.0)

    # Poll function should have been passed the result of fn
    assert_that(last_descriptors[0].result, equal_to(123))

    # It should be possible to cancel the future
    assert_that(future.cancel())

    # Wait until next poll
    poll_ran.clear()
    poll_ran.wait(10.0)

    # The cancelled future should have been removed from the
    # descriptors passed to the poll function
    assert_that(last_descriptors, equal_to([]))

    # It should be harmless to request cancel again
    assert_that(future.cancel())


def test_cancel_during_poll_fn(executor):
    queue = Queue()
    poll_ran = Event()
    last_descriptors = []

    def poll_fn(descriptors):
        while last_descriptors:
            last_descriptors.pop()
        last_descriptors.extend(descriptors)
        poll_ran.set()

        should_process = queue.get(True)
        if should_process:
            for descriptor in descriptors:
                if descriptor.result == "pass":
                    descriptor.yield_result("pass")
                else:
                    descriptor.yield_exception(RuntimeError("fail"))

    poll_executor = PollExecutor(executor, poll_fn, default_interval=0.01)
    futures = [poll_executor.submit(lambda x: x, x) for x in ("pass", "fail")]

    # Wait until both futures move to polling mode.
    def wait_two_futures():
        poll_ran.clear()
        queue.put(False)
        poll_ran.wait()
        assert_that(last_descriptors, has_length(2))

    assert_soon(wait_two_futures)

    # OK, now wait until the poll function is in the middle of executing
    poll_ran.clear()
    queue.put(False)
    poll_ran.wait()

    # Cancel the futures while poll function is in progress
    assert_that(futures[0].cancel())
    assert_that(futures[1].cancel())

    # Now let the poll function proceed, and attempt to update the futures
    poll_ran.clear()
    queue.put(True)
    poll_ran.wait()

    # The futures should remain cancelled
    assert_that(futures[0].cancelled())
    assert_that(futures[1].cancelled())

    # And they should not be passed to the poll function any more
    assert_that(last_descriptors, equal_to([]))


def test_poll_fail(executor):
    task_id_queue = Queue()
    tasks = {}
    poll_should_fail = [False]
    poll_ran = Event()

    last_descriptors = []

    def poll_fn(descriptors):
        while last_descriptors:
            last_descriptors.pop()
        last_descriptors.extend(descriptors)
        poll_ran.set()
        if poll_should_fail[0]:
            raise RuntimeError("simulated poll error")
        return poll_tasks(tasks, descriptors)

    poll_executor = PollExecutor(executor, poll_fn, default_interval=0.01)

    def make_task(x):
        return "%s-%s" % (x, task_id_queue.get(True))

    inputs = ["a", "b", "c"]
    futures = [poll_executor.submit(make_task, x) for x in inputs]

    # The futures should not currently be able to progress.
    assert_that(not any([f.done() for f in futures]))

    # Allow tasks to be created.
    task_id_queue.put("x")
    task_id_queue.put("y")
    task_id_queue.put("z")

    # Let one of the tasks complete
    tasks["a-x"] = "done"
    tasks["a-y"] = "done"
    tasks["a-z"] = "done"

    # Future a should become resolved
    assert_that(futures[0].result(10), equal_to("done"))

    # Now set up the poll function to fail
    poll_should_fail[0] = True

    # That should make both remaining futures fail
    assert_that(
        calling(futures[1].result).with_args(10),
        raises(RuntimeError, "simulated poll error"),
    )

    assert_that(
        calling(futures[2].result).with_args(10),
        raises(RuntimeError, "simulated poll error"),
    )

    # Wait for the next poll
    poll_ran.clear()
    poll_ran.wait(10)

    # The failed futures should no longer be passed into the poll function
    assert_that(last_descriptors, equal_to([]))