import asyncio
import os
import signal
import threading
import time as ttime
import types
from collections import defaultdict
from functools import partial
from traceback import FrameSummary, extract_tb

import pytest
from event_model import DocumentNames

from bluesky import Msg, RunEngine
from bluesky.plan_stubs import (
    abs_set,
    checkpoint,
    configure,
    declare_stream,
    pause,
    trigger_and_read,
    wait,
    wait_for,
)
from bluesky.plans import count, grid_scan, scan
from bluesky.preprocessors import (
    SupplementalData,
    baseline_wrapper,
    finalize_wrapper,
    reset_positions_decorator,
    rewindable_wrapper,
    run_decorator,
    run_wrapper,
    subs_wrapper,
)
from bluesky.protocols import Status
from bluesky.run_engine import (
    FailedStatus,
    IllegalMessageSequence,
    NoReplayAllowed,
    RequestAbort,
    RequestStop,
    RunEngineInterrupted,
    RunEngineStateMachine,
    TransitionError,
    WaitForTimeoutError,
)
from bluesky.tests import requires_ophyd, uses_os_kill_sigint
from bluesky.tests.utils import DocCollector, MsgCollector

from .utils import _careful_event_set, _fabricate_asycio_event


def test_states():
    assert RunEngineStateMachine.States.states() == [
        "idle",
        "running",
        "pausing",
        "paused",
        "halting",
        "stopping",
        "aborting",
        "suspending",
        "panicked",
    ]


def test_panic_trap(RE):
    RE._state = "panicked"
    for k in RunEngineStateMachine.States.states():
        if k != "panicked":
            with pytest.raises(TransitionError):
                RE._state = k


def test_state_is_readonly(RE):
    with pytest.raises(AttributeError):
        RE.state = "running"


@pytest.mark.parametrize(
    "deferred_pause_delay, is_pause_set",
    [
        (None, False),  # Do not pause
        (0.1, False),  # Pause before the checkpoint
        (1.4, True),  # Pause after the checkpoint
    ],
)
def test_deferred_pause_requested(RE, deferred_pause_delay, is_pause_set):
    """
    Test for ``deferred_pause_requested``.
    """
    assert RE.deferred_pause_requested is False

    plan = [Msg("sleep", None, 1), Msg("checkpoint"), Msg("sleep", None, 1)]

    pause_req_immediate = None

    def _pause():
        # This function is requesting the deferred pause and reads the property
        nonlocal pause_req_immediate
        RE.request_pause(defer=True)
        ttime.sleep(0.1)
        pause_req_immediate = RE.deferred_pause_requested

    t = None
    if deferred_pause_delay is not None:
        t = threading.Timer(deferred_pause_delay, _pause)
        t.start()

    try:
        RE(plan)
    except Exception:
        pass

    if deferred_pause_delay is not None:
        assert pause_req_immediate is True

    assert RE.deferred_pause_requested is is_pause_set
    if t is not None:
        t.join()


def test_verbose(RE, hw):
    RE.verbose = True
    assert RE.verbose
    # Emit all four kinds of document, exercising the logging.
    RE(
        [
            Msg("open_run"),
            Msg("declare_stream", None, hw.det, name="primary"),
            Msg("create", name="primary"),
            Msg("read", hw.det),
            Msg("save"),
            Msg("close_run"),
        ]
    )


def test_reset(RE):
    with pytest.raises(RunEngineInterrupted):
        RE([Msg("open_run"), Msg("pause")])

    assert len(RE._run_bundlers) > 0
    RE.reset()
    assert len(RE._run_bundlers) == 0


def test_running_from_paused_state_raises(RE):
    with pytest.raises(RunEngineInterrupted):
        RE([Msg("pause")])
    assert RE.state == "paused"
    with pytest.raises(RuntimeError):
        RE([Msg("null")])
    RE.resume()
    assert RE.state == "idle"
    RE([Msg("null")])


def test_resuming_from_idle_state_raises(RE):
    with pytest.raises(RuntimeError):
        RE.resume()
    with pytest.raises(RunEngineInterrupted):
        RE([Msg("pause")])
    assert RE.state == "paused"
    RE.resume()
    assert RE.state == "idle"
    with pytest.raises(RuntimeError):
        RE.resume()


def test_stopping_from_idle_state_raises(RE):
    with pytest.raises(TransitionError):
        RE.stop()


def test_pausing_from_idle_state_raises(RE):
    with pytest.raises(TransitionError):
        RE.request_pause(defer=False)


def test_aborting_from_idle_state_raises(RE):
    with pytest.raises(TransitionError):
        RE.abort()


def test_register(RE):
    mutable = {}
    RE.verbose = True

    async def func(msg):
        mutable["flag"] = True

    def plan():
        yield Msg("custom-command")

    RE.register_command("custom-command", func)
    RE(plan())
    assert "flag" in mutable
    # Unregister command; now the Msg should not be recognized.
    RE.unregister_command("custom-command")
    with pytest.raises(KeyError):
        RE([Msg("custom-command")])


def test_stop_motors_and_log_any_errors(RE, hw):
    # test that if stopping one motor raises an error, we can carry on
    stopped = {}

    def stop(self, *, success=False):
        stopped[self.name] = True

    def stop_encounters_error(self, *, success=False):
        stopped[self.name] = True
        raise Exception

    motor = hw.motor1
    broken_motor = hw.motor2
    motor.stop = types.MethodType(stop, motor)
    broken_motor.stop = types.MethodType(stop_encounters_error, broken_motor)

    with pytest.raises(RunEngineInterrupted):
        RE([Msg("set", broken_motor, 1), Msg("set", motor, 1), Msg("pause")])
    assert "motor1" in stopped
    assert "motor2" in stopped
    RE.stop()

    with pytest.raises(RunEngineInterrupted):
        RE([Msg("set", motor, 1), Msg("set", broken_motor, 1), Msg("pause")])
    assert "motor1" in stopped
    assert "motor2" in stopped
    RE.stop()


@requires_ophyd
def test_unstage_and_log_errors(RE):
    unstaged = {}

    from ophyd.sim import SynAxis

    class MoverWithFlag(SynAxis):
        def stage(self):
            return [self]

        def unstage(self):
            unstaged[self.name] = True
            return [self]

    class BrokenMoverWithFlag(SynAxis):
        def stage(self):
            return [self]

        def unstage(self):
            unstaged[self.name] = True
            return [self]

    a = MoverWithFlag(name="a")
    b = BrokenMoverWithFlag(name="b")

    unstaged.clear()
    RE([Msg("stage", a), Msg("stage", b)])
    assert "a" in unstaged
    assert "b" in unstaged

    unstaged.clear()
    RE([Msg("stage", b), Msg("stage", a)])
    assert "a" in unstaged
    assert "b" in unstaged


def test_open_run_twice_is_illegal(RE):
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("open_run"), Msg("open_run")])


def test_saving_without_an_open_bundle_is_illegal(RE):
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("open_run"), Msg("save")])


def test_dropping_without_an_open_bundle_is_illegal(RE):
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("open_run"), Msg("drop")])


def test_opening_a_bundle_without_a_run_is_illegal(RE):
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("create", name="primary")])


def test_checkpoint_inside_a_bundle_is_illegal(RE):
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("open_run"), Msg("create", name="primary"), Msg("checkpoint")])


def test_redundant_monitors_are_illegal(RE):
    class Dummy:
        def __init__(self, name):
            self.name = name

        def read_configuration(self):
            return {}

        def describe_configuration(self):
            return {}

        def describe(self):
            return {}

        def read(self):
            return {}

        def subscribe(self, *args, **kwargs):
            pass

        def clear_sub(self, *args, **kwargs):
            pass

    dummy = Dummy("dummy")

    with pytest.raises(IllegalMessageSequence):
        RE([Msg("open_run"), Msg("monitor", dummy), Msg("monitor", dummy)])

    # Monitoring, unmonitoring, and monitoring again is legal.
    RE([Msg("open_run"), Msg("monitor", dummy), Msg("unmonitor", dummy), Msg("monitor", dummy)])

    # Monitoring outside a run is illegal.
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("monitor", dummy)])

    # Unmonitoring something that was never monitored is illegal.
    with pytest.raises(IllegalMessageSequence):
        RE([Msg("unmonitor", dummy)])


def test_empty_bundle(RE, hw):
    mutable = {}

    def cb(name, doc):
        mutable["flag"] = True

    # In this case, an Event should be emitted.
    mutable.clear()
    RE(
        [
            Msg("open_run"),
            Msg("declare_stream", None, hw.det, name="primary"),
            Msg("create", name="primary"),
            Msg("read", hw.det),
            Msg("save"),
        ],
        {"event": cb},
    )
    assert "flag" in mutable

    # In this case, an Event should not be emitted because the bundle is
    # emtpy (i.e., there are no readings.)
    mutable.clear()
    RE(
        [Msg("open_run"), Msg("declare_stream", None, name="primary"), Msg("create", name="primary"), Msg("save")],
        {"event": cb},
    )
    assert "flag" not in mutable


def test_dispatcher_unsubscribe_all(RE):
    def count_callbacks(RE):
        return sum([len(cbs) for cbs in RE.dispatcher.cb_registry.callbacks.values()])

    def cb(name, doc):
        pass

    RE.subscribe(cb)
    assert count_callbacks(RE) == len(DocumentNames)
    RE.dispatcher.unsubscribe_all()
    assert count_callbacks(RE) == 0


def test_stage_and_unstage_are_optional_methods(RE):
    class Dummy:
        pass

    dummy = Dummy()

    RE([Msg("stage", dummy), Msg("unstage", dummy)])


def test_need_both_stage_and_unstage_for_stagable(RE):
    class StageOnly:
        def stage(self):
            raise Exception

    class UnstageOnly:
        def unstage(self):
            raise Exception

    stage_only = StageOnly()
    unstage_only = UnstageOnly()
    RE(
        [
            Msg("stage", stage_only),
            Msg("stage", unstage_only),
            Msg("unstage", stage_only),
            Msg("unstage", unstage_only),
        ]
    )


def test_pause_resume_devices(RE):
    paused = {}
    resumed = {}

    class Dummy:
        def __init__(self, name):
            self.name = name

        def pause(self):
            paused[self.name] = True

        def resume(self):
            resumed[self.name] = True

    dummy = Dummy("dummy")

    with pytest.raises(RunEngineInterrupted):
        RE([Msg("stage", dummy), Msg("pause")])
    RE.resume()
    assert "dummy" in paused
    assert "dummy" in resumed


def test_stage_and_unstage_status_objects(RE):
    from ophyd import StatusBase

    staged = {}
    unstaged = {}
    started_stage = {}
    started_unstage = {}

    class Dummy:
        def __init__(self, name: str) -> None:
            self.name: str = name

        def _callback(self, d: dict, st: Status):
            d[self.name] = True
            st.set_finished()  # type: ignore

        def stage(self) -> Status:
            st = StatusBase()
            started_stage[self.name] = True
            threading.Timer(1, self._callback, args=[staged, st]).start()
            return st

        def unstage(self) -> Status:
            st = StatusBase()
            started_unstage[self.name] = True
            threading.Timer(1, self._callback, args=[unstaged, st]).start()
            return st

    dummy1 = Dummy("d1")
    dummy2 = Dummy("d2")

    def my_plan():
        yield Msg("stage", dummy1, group="stage_group")
        yield Msg("stage", dummy2, group="stage_group")

        assert "d1" in started_stage and "d2" in started_stage
        assert not staged

        yield Msg("wait", group="stage_group")
        assert "d1" in staged and "d2" in staged

        yield Msg("unstage", dummy1, group="unstage_group")
        yield Msg("unstage", dummy2, group="unstage_group")

        assert "d1" in started_unstage and "d2" in started_unstage
        assert not unstaged

        yield Msg("wait", group="unstage_group")
        assert "d1" in unstaged and "d2" in unstaged

    start = ttime.monotonic()
    RE(my_plan())
    stop = ttime.monotonic()

    assert 2 <= stop - start < 3


def test_bad_call_args(RE):
    with pytest.raises(Exception):  # noqa: B017
        RE(53)
    assert RE.state == "idle"


def test_record_interruptions(RE):
    docs = defaultdict(lambda: [])

    def collect(name, doc):
        print("HI", name)
        docs[name].append(doc)
        print(docs)

    RE.subscribe(collect)
    RE.ignore_callback_exceptions = False
    RE.msg_hook = print

    # The 'pause' inside the run should generate an event iff
    # record_interruptions is True.
    plan = [Msg("pause"), Msg("open_run"), Msg("pause"), Msg("close_run")]

    assert not RE.record_interruptions
    with pytest.raises(RunEngineInterrupted):
        RE(plan)
    with pytest.raises(RunEngineInterrupted):
        RE.resume()
    RE.resume()
    assert len(docs["descriptor"]) == 0
    assert len(docs["event"]) == 0

    RE.record_interruptions = True
    with pytest.raises(RunEngineInterrupted):
        RE(plan)
    with pytest.raises(RunEngineInterrupted):
        RE.resume()
    RE.resume()
    assert len(docs["descriptor"]) == 1
    assert len(docs["event"]) == 2
    docs["event"][0]["data"]["interruption"] == "pause"  # noqa: B015
    docs["event"][1]["data"]["interruption"] == "resume"  # noqa: B015


@requires_ophyd
def _make_unrewindable_marker():
    from ophyd.sim import SynAxis, SynGauss

    class UnReplayableSynGauss(SynGauss):
        def pause(self):
            raise NoReplayAllowed()

    motor = SynAxis(name="motor")

    def test_plan(motor, det):
        yield Msg("set", motor, 0)
        yield Msg("trigger", det)
        yield Msg("pause")
        yield Msg("set", motor, 1)
        yield Msg("trigger", det)

    inps = []
    inps.append(
        (
            test_plan,
            motor,
            UnReplayableSynGauss("det", motor, "motor", center=0, Imax=1),
            ["set", "trigger", "pause", "set", "trigger"],
        )
    )

    inps.append(
        (
            test_plan,
            motor,
            SynGauss("det", motor, "motor", center=0, Imax=1),
            ["set", "trigger", "pause", "set", "trigger", "set", "trigger"],
        )
    )

    return pytest.mark.parametrize("plan,motor,det,msg_seq", inps)


@_make_unrewindable_marker()
def test_unrewindable_det(RE, plan, motor, det, msg_seq):
    msgs = []

    def collector(msg):
        msgs.append(msg)

    RE.msg_hook = collector
    with pytest.raises(RunEngineInterrupted):
        RE(plan(motor, det))
    RE.resume()
    assert [m.command for m in msgs] == msg_seq


@requires_ophyd
def _make_unrewindable_suspender_marker():
    from ophyd.sim import SynAxis, SynGauss

    class UnReplayableSynGauss(SynGauss):
        def pause(self):
            raise NoReplayAllowed()

    motor = SynAxis(name="motor")

    def test_plan(motor, det):
        yield Msg("set", motor, 0)
        yield Msg("trigger", det)
        yield Msg("sleep", None, 1)
        yield Msg("set", motor, 0)
        yield Msg("trigger", det)

    inps = []
    inps.append(
        (
            test_plan,
            motor,
            UnReplayableSynGauss("det", motor, "motor", center=0, Imax=1),
            [
                "set",
                "trigger",
                "sleep",
                "_start_suspender",
                "rewindable",
                "wait_for",
                "_resume_from_suspender",
                "rewindable",
                "set",
                "trigger",
            ],
        )
    )

    inps.append(
        (
            test_plan,
            motor,
            SynGauss("det", motor, "motor", center=0, Imax=1),
            [
                "set",
                "trigger",
                "sleep",
                "_start_suspender",
                "rewindable",
                "wait_for",
                "_resume_from_suspender",
                "rewindable",
                "set",
                "trigger",
                "sleep",
                "set",
                "trigger",
            ],
        )
    )

    return pytest.mark.parametrize("plan,motor,det,msg_seq", inps)


@_make_unrewindable_suspender_marker()
def test_unrewindable_det_suspend(RE, plan, motor, det, msg_seq):
    from bluesky.utils import ts_msg_hook

    msgs = []
    loop = RE.loop

    def collector(msg):
        ts_msg_hook(msg)
        msgs.append(msg)

    RE.msg_hook = collector

    ev = _fabricate_asycio_event(loop)

    timer = threading.Timer(0.5, RE.request_suspend, kwargs=dict(fut=ev.wait))  # noqa: C408
    timer.start()

    def verbose_set():
        print("seting")
        ev.set()

    loop.call_soon_threadsafe(loop.call_later, 1, verbose_set)

    RE(plan(motor, det))
    assert [m.command for m in msgs] == msg_seq
    timer.join()


@pytest.mark.parametrize("unpause_func", [lambda RE: RE.stop(), lambda RE: RE.abort(), lambda RE: RE.resume()])
def test_cleanup_after_pause(RE, unpause_func, hw):
    motor = hw.motor
    motor.set(1024)

    @reset_positions_decorator()
    def simple_plan(motor):
        for j in range(15):
            yield Msg("set", motor, j)
        yield Msg("pause")
        for j in range(15):
            yield Msg("set", motor, -j)

    with pytest.raises(RunEngineInterrupted):
        RE(simple_plan(motor))
    assert motor.position == 14
    unpause_func(RE)
    assert motor.position == 1024


@pytest.mark.parametrize(
    "unpause_func,excp",
    [
        (lambda RE: RE.stop(), RequestStop),
        (lambda RE: RE.abort(), RequestAbort),
        (lambda RE: RE.halt(), GeneratorExit),
    ],
)
def test_exit_raise(RE, unpause_func, excp):
    flag = False

    @reset_positions_decorator()
    def simple_plan():
        nonlocal flag
        try:
            yield from checkpoint()
            yield Msg("pause")
        except excp:
            flag = True

    with pytest.raises(RunEngineInterrupted):
        RE(simple_plan())
    unpause_func(RE)
    assert flag


@uses_os_kill_sigint
def test_sigint_three_hits(RE, hw):
    import time

    motor = hw.motor
    motor.delay = 0.5

    pid = os.getpid()

    def sim_kill(n):
        for j in range(n):  # noqa: B007
            time.sleep(0.05)
            os.kill(pid, signal.SIGINT)

    lp = RE.loop
    motor.loop = lp

    def self_sig_int_plan():
        threading.Timer(0.05, sim_kill, (3,)).start()
        yield from abs_set(motor, 1, wait=True)

    start_time = ttime.time()
    with pytest.raises(RunEngineInterrupted):
        RE(finalize_wrapper(self_sig_int_plan(), abs_set(motor, 0, wait=True)))
    end_time = ttime.time()
    # not enough time for motor to cleanup, but long enough to start
    assert 0.05 < end_time - start_time < 0.2
    RE.abort()  # now cleanup

    done_cleanup_time = ttime.time()
    # this should be 0.5 (the motor.delay) above, leave sloppy for CI
    assert 0.3 < done_cleanup_time - end_time < 0.6


@uses_os_kill_sigint
def test_sigint_many_hits_pln(RE):
    pid = os.getpid()

    def sim_kill(n):
        for j in range(n):
            print("KILL", j)
            ttime.sleep(0.05)
            os.kill(pid, signal.SIGINT)

    def hanging_plan():
        "a plan that blocks the RunEngine's normal Ctrl+C handing with a sleep"
        for j in range(100):  # noqa: B007
            ttime.sleep(0.1)
        yield Msg("null")

    start_time = ttime.time()
    timer = threading.Timer(0.2, sim_kill, (11,))
    timer.start()
    with pytest.raises(RunEngineInterrupted):
        RE(hanging_plan())
    # Check that hammering SIGINT escaped from that 10-second sleep.
    assert ttime.time() - start_time < 5
    # The KeyboardInterrupt will have been converted to a hard pause that
    # the test plan can not handle so we abort and go to idle.
    assert RE.state == "idle"
    timer.join()


@uses_os_kill_sigint
def test_sigint_many_hits_panic(RE):
    raise pytest.skip("hangs tests on exit")
    pid = os.getpid()

    def sim_kill(n):
        for j in range(n):
            print("KILL", j, ttime.monotonic() - start_time)
            ttime.sleep(0.05)
            os.kill(pid, signal.SIGINT)

    def hanging_plan():
        "a plan that blocks the RunEngine's normal Ctrl+C handing with a sleep"
        yield Msg("null")
        ttime.sleep(5)
        yield Msg("null")

    start_time = ttime.monotonic()
    timer = threading.Timer(0.2, sim_kill, (11,))
    timer.start()
    with pytest.raises(RunEngineInterrupted):
        RE(hanging_plan())
    # Check that hammering SIGINT escaped from that 5-second sleep.
    assert (ttime.monotonic() - start_time) < 2.5
    # The KeyboardInterrupt but because we could not shut down, panic!
    assert RE.state == "panicked"
    timer.join()

    with pytest.raises(RuntimeError):
        RE([])

    with pytest.raises(RuntimeError):
        RE.stop()

    with pytest.raises(RuntimeError):
        RE.halt()

    with pytest.raises(RuntimeError):
        RE.abort()

    with pytest.raises(RuntimeError):
        RE.resume()

    with pytest.raises(RuntimeError):
        RE.request_pause()


@uses_os_kill_sigint
def test_sigint_many_hits_cb(RE):
    pid = os.getpid()

    def sim_kill(n):
        for j in range(n):  # noqa: B007
            print("KILL")
            ttime.sleep(0.05)
            os.kill(pid, signal.SIGINT)

    @run_decorator()
    def infinite_plan():
        while True:
            yield Msg("null")

    def hanging_callback(name, doc):
        for j in range(100):  # noqa: B007
            ttime.sleep(0.1)

    start_time = ttime.time()
    timer = threading.Timer(0.2, sim_kill, (11,))
    timer.start()
    with pytest.raises(RunEngineInterrupted):
        RE(infinite_plan(), {"start": hanging_callback})
    # Check that hammering SIGINT escaped from that 10-second sleep.
    assert ttime.time() - start_time < 5
    # The KeyboardInterrupt will have been converted to a hard pause.
    assert RE.state == "idle"
    timer.join()


@uses_os_kill_sigint
def test_no_context_manager(RE):
    # Clear the context managers so that RE will not react to sigint
    RE.context_managers = []

    # Proceed as normal
    pid = os.getpid()

    def sim_kill(n=1):
        for j in range(n):
            print("KILL", j)
            ttime.sleep(0.05)
            os.kill(pid, signal.SIGINT)

    def hanging_plan():
        "a plan that blocks the RunEngine's normal Ctrl+C handing with a sleep"
        ttime.sleep(2)
        yield Msg("null")

    # Only send one SIGINT
    timer = threading.Timer(0.5, sim_kill, (1,))
    timer.start()

    t = threading.Thread(target=RE, args=(hanging_plan()))
    start = ttime.time()
    t.start()

    # Wait for the KeyboardInterrupt and handle in the main thread
    with pytest.raises(KeyboardInterrupt):
        ttime.sleep(5)

    t.join()
    delta = ttime.time() - start

    # Hanging plan finished, but extra sleep did not
    assert 2 < delta < 5
    timer.join()


def test_many_context_managers(RE):
    class Manager:
        enters = 0
        exits = 0

        def __init__(self, RE):
            pass

        def __enter__(self):
            Manager.enters += 1

        def __exit__(self, *args, **kwargs):
            Manager.exits += 1

    n = 42
    RE.context_managers.extend([Manager] * n)
    RE([Msg("null")])
    assert Manager.enters == n
    assert Manager.exits == n


def _make_plan_marker():
    @reset_positions_decorator()
    def raiser(motor):
        for j in range(15):
            yield Msg("set", motor, j)
        raise RuntimeError()

    @reset_positions_decorator()
    def pausing_raiser(motor):
        for j in range(15):
            yield Msg("set", motor, j)
        yield Msg("pause")
        raise RuntimeError()

    @reset_positions_decorator()
    def bad_set(motor):
        for j in range(15):
            yield Msg("set", motor, j)
            yield Msg("set", None, j)

    @reset_positions_decorator()
    def bad_msg(motor):
        for j in range(15):
            yield Msg("set", motor, j)
        yield Msg("aardvark")

    @reset_positions_decorator()
    def cannot_pauser(motor):
        yield Msg("clear_checkpoint")
        for j in range(15):
            yield Msg("set", motor, j)
        yield Msg("pause")

    return pytest.mark.parametrize("plan", [raiser, bad_set, bad_msg, pausing_raiser, cannot_pauser])


@_make_plan_marker()
def test_cleanup_pathological_plans(RE, hw, plan):
    motor = hw.motor
    motor.set(1024)
    try:
        try:
            RE(plan(motor))
        except RunEngineInterrupted:
            pass
        if RE.state == "paused":
            assert motor.position != 1024
            RE.resume()
    except Exception:
        pass
    assert motor.position == 1024


def test_finalizer_closeable():
    pre = (j for j in range(18))
    post = (j for j in range(18))

    plan = finalize_wrapper(pre, post)

    for j in range(3):  # noqa: B007
        next(plan)
    plan.close()


def test_invalid_generator(RE, hw, capsys):
    motor = hw.motor
    from bluesky.utils import ts_msg_hook

    RE.msg_hook = ts_msg_hook

    # this is not a valid generator as it will try to yield if it
    # is throw a GeneratorExit
    def patho_finalize_wrapper(plan, post):
        try:
            yield from plan
        finally:
            yield from post

    def base_plan(motor):
        for j in range(5):
            yield Msg("set", motor, j * 2 + 1)
        yield Msg("pause")

    def post_plan(motor):
        yield Msg("set", motor, 500)

    def pre_suspend_plan():
        yield Msg("set", motor, -500)
        raise GeneratorExit("this one")

    def make_plan():
        return patho_finalize_wrapper(base_plan(motor), post_plan(motor))

    with pytest.raises(RunEngineInterrupted):
        RE(make_plan())
    RE.request_suspend(None, pre_plan=pre_suspend_plan())
    capsys.readouterr()
    try:
        RE.resume()
    except ValueError as sf:
        assert sf.__cause__.args[0] == "this one"

    actual_err, _ = capsys.readouterr()
    expected_prefix = "The plan "
    expected_postfix = (" tried to yield a value on close.  Please fix your plan.\n")[::-1]
    assert actual_err[: len(expected_prefix)] == expected_prefix
    assert actual_err[::-1][: len(expected_postfix)] == expected_postfix


def test_exception_cascade_REside(RE):
    except_hit = False

    def pausing_plan():
        nonlocal except_hit
        for j in range(5):  # noqa: B007
            yield Msg("null")
        try:
            yield Msg("pause")
        except Exception:
            except_hit = True
            raise

    def pre_plan():
        yield Msg("aardvark")

    def post_plan():
        for j in range(5):  # noqa: B007
            yield Msg("null")

    with pytest.raises(RunEngineInterrupted):
        RE(pausing_plan())
    ev = _fabricate_asycio_event(RE.loop)
    ev.set()
    RE.request_suspend(ev.wait, pre_plan=pre_plan())
    with pytest.raises(KeyError):
        RE.resume()
    assert except_hit


def test_exception_cascade_planside(RE):
    except_hit = False

    def pausing_plan():
        nonlocal except_hit
        for j in range(5):  # noqa: B007
            yield Msg("null")
        try:
            yield Msg("pause")
        except Exception:
            except_hit = True
            raise

    def pre_plan():
        yield Msg("null")
        raise RuntimeError()

    def post_plan():
        for j in range(5):  # noqa: B007
            yield Msg("null")

    with pytest.raises(RunEngineInterrupted):
        RE(pausing_plan())
    ev = _fabricate_asycio_event(RE.loop)
    ev.set()
    RE.request_suspend(ev.wait, pre_plan=pre_plan())
    with pytest.raises(RuntimeError):
        RE.resume()
    assert except_hit


def test_sideband_cancel(RE):
    loop = RE.loop
    ev = _fabricate_asycio_event(RE.loop)

    def done():
        _careful_event_set(ev)()

    def side_band_kill():
        RE.loop.call_soon_threadsafe(RE._task.cancel)

    scan = [
        Msg(
            "wait_for",
            None,
            [
                ev.wait,
            ],
        ),
    ]
    assert RE.state == "idle"
    start = ttime.time()
    timer = threading.Timer(0.5, side_band_kill)
    timer.start()
    loop.call_soon_threadsafe(loop.call_later, 2, done)
    RE(scan)
    assert RE.state == "idle"
    assert RE._task.cancelled()
    stop = ttime.time()

    assert 0.5 < (stop - start) < 2
    timer.join()


def test_no_rewind(RE):
    msg_lst = []

    def msg_collector(msg):
        msg_lst.append(msg)

    RE.rewindable = False

    plan = [Msg("null")] * 3 + [Msg("pause")] + [Msg("null")] * 3
    RE.msg_hook = msg_collector
    with pytest.raises(RunEngineInterrupted):
        RE(plan)
    RE.resume()
    assert msg_lst == plan


def test_no_rewindable_msg(RE):
    RE.rewindable = True

    msg_lst = []

    def msg_collector(msg):
        msg_lst.append(msg)

    plan = [Msg("null")] * 3 + [Msg("pause"), Msg("rewindable", None, False)] + [Msg("null")] * 3

    RE.msg_hook = msg_collector

    with pytest.raises(RunEngineInterrupted):
        RE(plan)
    RE.resume()
    assert msg_lst[:4] == plan[:4]
    assert msg_lst[4:7] == plan[:3]
    assert msg_lst[7:] == plan[4:]


@pytest.mark.parametrize("start_state", [True, False])
def test_rewindable_state_retrival(RE, start_state):
    RE.rewindable = start_state

    def rewind_plan(start_value):
        ret = yield Msg("rewindable", None, None)
        assert ret is start_state
        cache_state = ret
        ret = yield Msg("rewindable", None, start_state)
        assert ret is start_state

        ret = yield Msg("rewindable", None, not start_state)
        assert ret is (not start_state)

        ret = yield Msg("rewindable", None, cache_state)
        assert ret is start_state

    RE(rewind_plan(start_state))
    assert RE.rewindable is start_state


@pytest.mark.parametrize(
    "start_state,msg_seq",
    (
        (
            True,
            [
                "open_run",
                "declare_stream",
                "rewindable",
                "rewindable",
                "trigger",
                "trigger",
                "wait",
                "create",
                "read",
                "read",
                "save",
                "rewindable",
                "close_run",
            ],
        ),
        (
            False,
            [
                "open_run",
                "declare_stream",
                "rewindable",
                "trigger",
                "trigger",
                "wait",
                "create",
                "read",
                "read",
                "save",
                "close_run",
            ],
        ),
    ),
)
def test_nonrewindable_detector(RE, hw, start_state, msg_seq):
    class FakeSig:
        def get(self):
            return False

    hw.det.rewindable = FakeSig()

    RE.rewindable = start_state
    m_col = MsgCollector()
    RE.msg_hook = m_col

    def test():
        yield from declare_stream(hw.motor, hw.det, name="primary")
        return (yield from trigger_and_read([hw.motor, hw.det]))

    RE(run_wrapper(test()))

    assert [m.command for m in m_col.msgs] == msg_seq


@pytest.mark.parametrize(
    "start_state,msg_seq",
    ((True, ["rewindable", "rewindable", "aardvark", "rewindable"]), (False, ["rewindable", "aardvark"])),
)
def test_nonrewindable_finalizer(RE, hw, start_state, msg_seq):
    class FakeSig:
        def get(self):
            return False

    det = hw.det
    det.rewindable = FakeSig()

    RE.rewindable = start_state
    m_col = MsgCollector()
    RE.msg_hook = m_col

    def evil_plan():
        assert RE.rewindable is False
        yield Msg("aardvark")

    with pytest.raises(KeyError):
        RE(rewindable_wrapper(evil_plan(), False))

    assert RE.rewindable is start_state

    assert [m.command for m in m_col.msgs] == msg_seq


def test_halt_from_pause(RE):
    except_hit = False
    m_coll = MsgCollector()
    RE.msg_hook = m_coll

    def pausing_plan():
        nonlocal except_hit
        for j in range(5):  # noqa: B007
            yield Msg("null")
        try:
            yield Msg("pause")
        except Exception:
            yield Msg("null")
            except_hit = True
            raise

    with pytest.raises(RunEngineInterrupted):
        RE(pausing_plan())
    RE.halt()
    assert not except_hit
    assert [m.command for m in m_coll.msgs] == ["null"] * 5 + ["pause"]


def test_halt_async(RE):
    except_hit = False
    m_coll = MsgCollector()
    RE.msg_hook = m_coll

    def sleeping_plan():
        nonlocal except_hit
        try:
            yield Msg("sleep", None, 50)
        except Exception:
            yield Msg("null")
            except_hit = True
            raise

    timer = threading.Timer(0.1, RE.halt)
    timer.start()
    start = ttime.time()
    with pytest.raises(RunEngineInterrupted):
        RE(sleeping_plan())
    stop = ttime.time()
    assert 0.09 < stop - start < 5
    assert not except_hit
    assert [m.command for m in m_coll.msgs] == ["sleep"]
    timer.join()


@pytest.mark.parametrize(
    "cancel_func", [lambda RE: RE.stop(), lambda RE: RE.abort(), lambda RE: RE.request_pause(defer=False)]
)
def test_prompt_stop(RE, cancel_func):
    except_hit = False
    m_coll = MsgCollector()
    RE.msg_hook = m_coll

    def sleeping_plan():
        nonlocal except_hit
        try:
            yield Msg("sleep", None, 50)
        except Exception:
            yield Msg("null")
            except_hit = True
            raise

    timer = threading.Timer(0.1, partial(cancel_func, RE))
    timer.start()
    start = ttime.time()
    with pytest.raises(RunEngineInterrupted):
        RE(sleeping_plan())
    stop = ttime.time()
    if RE.state != "idle":
        RE.abort()
    assert 0.09 < stop - start < 5
    assert except_hit
    assert [m.command for m in m_coll.msgs] == ["sleep", "null"]
    timer.join()


@pytest.mark.parametrize(
    "change_func",
    [
        lambda RE: RE.stop(),
        lambda RE: RE.abort(),
        lambda RE: RE.halt(),
        lambda RE: RE.request_pause(),
        lambda RE: RE.request_pause(defer=True),
        lambda RE: RE.resume(),
    ],
)
def test_bad_from_idle_transitions(RE, change_func):
    with pytest.raises(TransitionError):
        change_func(RE)


def test_empty_cache_pause(RE):
    RE.rewindable = False
    pln = [
        Msg("open_run"),
        Msg("declare_stream", None, name="primary"),
        Msg("create", name="primary"),
        Msg("pause"),
        Msg("save"),
        Msg("close_run"),
    ]
    with pytest.raises(RunEngineInterrupted):
        RE(pln)
    RE.resume()


def test_state_hook(RE):
    states = []

    def log_state(new, old):
        states.append((new, old))

    RE.state_hook = log_state
    with pytest.raises(RunEngineInterrupted):
        RE([Msg("open_run"), Msg("pause"), Msg("close_run")])
    RE.resume()
    expected = [
        ("running", "idle"),
        ("pausing", "running"),
        ("paused", "pausing"),
        ("running", "paused"),
        ("idle", "running"),
    ]
    assert states == expected


def test_max_depth(RE):
    RE.max_depth is None  # noqa: B015
    RE([])  # should not raise

    # assume test framework needs less than 100 stacks... haha
    RE.max_depth = 100
    RE([])  # should not raise

    RE.max_depth = 0
    with pytest.raises(RuntimeError):
        RE([])


def test_preprocessors(RE):
    def custom_cleanup(plan):
        yield from plan
        yield Msg("null", "cleanup")  # just a sentinel

    def my_sub(name, doc):
        pass

    def custom_subs(plan):
        yield from subs_wrapper(plan, my_sub)

    RE.preprocessors = [custom_cleanup, custom_subs]
    actual = []
    RE.msg_hook = lambda msg: actual.append(msg)
    RE([Msg("null")])
    print(actual)
    expected = [
        Msg("subscribe", None, my_sub, "all"),
        Msg("null"),
        Msg("null", "cleanup"),
        Msg("unsubscribe", None, token=0),
    ]
    assert actual == expected


@requires_ophyd
def test_pardon_failures(RE):
    from ophyd import StatusBase

    st = StatusBase()

    class Dummy:
        name = "dummy"

        def set(self, val):
            return st

    dummy = Dummy()

    RE([Msg("set", dummy, 1)])
    st._finished(success=False)
    RE([Msg("null")])


@requires_ophyd
def test_exceptions_kill_run_and_failed_status_holds_detail(RE):
    # just to make sure that 'pardon_failures' does not block *real* failures
    from ophyd import StatusBase
    from ophyd.utils.errors import UnknownStatusFailure

    class Dummy:
        name = "dummy"

        def set(self, val):
            st = StatusBase()
            st._finished(success=False)
            return st

    dummy = Dummy()

    with pytest.raises(FailedStatus) as exc:
        RE([Msg("set", dummy, 1, group="test"), Msg("wait", group="test")])
        assert isinstance(exc.args[0], UnknownStatusFailure)


@requires_ophyd
def test_wait_with_timeout_can_be_repeated(RE):
    from ophyd import StatusBase

    class Dummy:
        name = "dummy"

        def set(self, val):
            self.st = StatusBase()
            return self.st

    def plan():
        dummy = Dummy()
        yield Msg("set", dummy, 1, group="test")
        with pytest.raises(TimeoutError):
            yield Msg("wait", group="test", timeout=0.1)
        with pytest.raises(TimeoutError):
            yield Msg("wait", group="test", timeout=0.1)
        dummy.st.set_finished()
        yield Msg("wait", group="test")

    RE(plan())


@requires_ophyd
def test_status_propagates_exception_through_run_engine(RE: RunEngine):
    # just to make sure that 'pardon_failures' does not block *real* failures
    from ophyd import StatusBase

    class DummyV1:
        name = "dummy"

        def set(self, val):
            st = StatusBase()
            try:
                1 / 0  # noqa: B018
            except ZeroDivisionError as exc:
                st.set_exception(exc)

            return st

    dummy1 = DummyV1()

    with pytest.raises(FailedStatus) as exc:
        RE([Msg("set", dummy1, 1, group="test"), Msg("wait", group="test")])

        traceback: list[FrameSummary] = extract_tb(exc.__traceback__)  # type: ignore[attr-defined]
        assert traceback[0].filename == __file__
        assert traceback[0].line == "RE([Msg('set', dummy1, 1, group='test'),"
        assert traceback[-1].name == "set"
        assert traceback[-1].line == "1/0"

        assert isinstance(exc.args[0], ZeroDivisionError)  # type: ignore[attr-defined]


def test_colliding_streams(RE, hw):
    collector = {"primary": [], "baseline": []}
    descs = {}

    def local_cb(name, doc):
        if name == "descriptor":
            descs[doc["uid"]] = doc["name"]
        elif name == "event":
            collector[descs[doc["descriptor"]]].append(doc)

    RE(
        baseline_wrapper(
            grid_scan([hw.motor], hw.motor, -1, 1, 5, hw.motor1, -5, 5, 7, True), [hw.motor, hw.motor1]
        ),
        local_cb,
    )

    assert len(collector["primary"]) == 35
    assert len(collector["baseline"]) == 2

    assert list(range(1, 36)) == [e["seq_num"] for e in collector["primary"]]
    assert list(range(1, 3)) == [e["seq_num"] for e in collector["baseline"]]


def test_old_subscribe(RE):
    # Old usage had reversed argument order. It should warn but still work.
    collector = []

    def collect(name, doc):
        collector.append(doc)

    with pytest.warns(UserWarning):
        RE.subscribe("all", collect)

    RE([Msg("open_run"), Msg("close_run")])
    assert len(collector) == 2

    RE.unsubscribe(0)
    with pytest.warns(UserWarning):
        RE.subscribe("start", collect)

    RE([Msg("open_run"), Msg("close_run")])
    assert len(collector) == 3

    RE.unsubscribe(1)


def test_waiting_hook(RE, hw):
    collector = []

    def collect(sts):
        collector.append(sts)

    RE.waiting_hook = collect

    RE([Msg("set", hw.motor, 5, group="A"), Msg("wait", group="A")])

    sts, none = collector
    assert isinstance(sts, set)
    assert len(sts) == 1
    assert none is None
    collector.clear()

    RE([Msg("set", hw.motor1, 5, group="A"), Msg("set", hw.motor2, 3, group="A"), Msg("wait", group="A")])

    sts, none = collector
    assert isinstance(sts, set)
    assert len(sts) == 2
    assert none is None
    collector.clear()


def test_hints(RE):
    class Detector:
        def __init__(self, name):
            self.name = name
            self.parent = None
            self.hints = {"vis": "placeholder"}

        async def read(self):
            await asyncio.sleep(0.1)
            return {}

        async def describe(self):
            return {}

        async def read_configuration(self):
            return {}

        async def describe_configuration(self):
            return {}

    det = Detector("det")

    collector = []

    RE(count([det]), {"descriptor": lambda name, doc: collector.append(doc)})
    doc = collector.pop()
    assert doc["hints"]["det"] == {"vis": "placeholder"}


def test_filled(RE, hw, db):
    collector = []

    def collect(name, doc):
        if name == "event":
            collector.append(doc)

    RE(count([hw.det]), collect)

    (event,) = collector
    assert event["filled"] == {}
    collector.clear()

    RE(count([hw.img]), collect)
    (event,) = collector
    assert event["filled"] == {"img": False}


def test_double_call(RE):
    uid1 = RE(count([]))
    uid2 = RE(count([]))

    assert uid1 != uid2


@pytest.mark.parametrize("predeclare", [True, False])
def test_num_events(RE, hw, db, predeclare, monkeypatch):
    if predeclare:
        monkeypatch.setenv("BLUESKY_PREDECLARE", "1")
    RE.subscribe(db.insert)

    rs1 = RE(count([]))
    if RE.call_returns_result:
        uid1 = rs1.run_start_uids[0]
    else:
        uid1 = rs1[0]
    h = db[uid1]
    if predeclare:
        assert h.stop["num_events"] == {"primary": 0}
    else:
        assert h.stop["num_events"] == {}

    rs2 = RE(count([hw.det], 5))
    if RE.call_returns_result:
        uid2 = rs2.run_start_uids[0]
    else:
        uid2 = rs2[0]
    h = db[uid2]
    assert h.stop["num_events"] == {"primary": 5}

    sd = SupplementalData(baseline=[hw.det])
    RE.preprocessors.append(sd)

    rs3 = RE(count([]))
    if RE.call_returns_result:
        uid3 = rs3.run_start_uids[0]
    else:
        uid3 = rs3[0]
    h = db[uid3]
    if predeclare:
        assert h.stop["num_events"] == {"primary": 0, "baseline": 2}
    else:
        assert h.stop["num_events"] == {"baseline": 2}

    rs4 = RE(count([hw.det], 5))
    if RE.call_returns_result:
        uid4 = rs4.run_start_uids[0]
    else:
        uid4 = rs4[0]
    h = db[uid4]
    assert h.stop["num_events"] == {"primary": 5, "baseline": 2}


def test_raise_if_interrupted_deprecation(RE):
    with pytest.warns(UserWarning):
        RE([], raise_if_interrupted=True)


@pytest.mark.parametrize(
    "bail_func,status", (("resume", "success"), ("stop", "success"), ("abort", "abort"), ("halt", "abort"))
)
def test_force_stop_exit_status(bail_func, status, RE):
    d = DocCollector()
    RE.subscribe(d.insert)

    @run_decorator()
    def bad_plan():
        print("going in")
        yield Msg("pause")

    with pytest.raises(RunEngineInterrupted):
        RE(bad_plan())

    rs = getattr(RE, bail_func)()

    if RE.call_returns_result:
        if bail_func == "resume":
            assert rs.plan_result is not RE.NO_PLAN_RETURN
        else:
            assert rs.plan_result is RE.NO_PLAN_RETURN
        uid = rs.run_start_uids[0]
        assert rs.exit_status == status
    else:
        uid = rs[0]
    assert len(d.start) == 1
    assert d.start[0]["uid"] == uid
    assert len(d.stop) == 1
    assert d.stop[uid]["exit_status"] == status


def test_exceptions_exit_status(RE):
    d = DocCollector()
    RE.subscribe(d.insert)

    class Snowflake(Exception): ...

    @run_decorator()
    def bad_plan():
        yield Msg("null")
        raise Snowflake("boo")

    with pytest.raises(Snowflake) as sf:
        RE(bad_plan())

    assert len(d.start) == 1
    rs = d.start[0]["uid"]
    assert len(d.stop) == 1
    assert d.stop[rs]["exit_status"] == "fail"
    assert d.stop[rs]["reason"] == str(sf.value)


def test_plan_return(RE):
    if not RE.call_returns_result:
        pytest.skip()

    def test_plan():
        yield Msg("null")
        return "plan_result"

    rs = RE(test_plan())
    assert rs.plan_result == "plan_result"
    assert rs.exit_status == "success"
    assert rs.interrupted is False


def test_plan_return_resume(RE):
    if not RE.call_returns_result:
        pytest.skip()

    def test_plan():
        yield Msg("null")
        yield Msg("pause")
        return "plan_result"

    with pytest.raises(RunEngineInterrupted):
        RE(test_plan())
    rs = RE.resume()
    assert rs.plan_result == "plan_result"
    assert rs.exit_status == "success"


def test_drop(RE, hw):
    det = hw.det

    def inner(msg):
        yield Msg("create", name="primary")
        yield Msg("read", det)
        yield Msg(msg)

    # Drop first, drop after saving, save after dropping
    def plan():
        yield Msg("open_run")
        yield Msg("declare_stream", None, det, name="primary")
        yield from inner("drop")
        yield from inner("save")
        yield from inner("drop")
        yield from inner("save")
        yield from inner("save")
        yield Msg("close_run")

    docs = defaultdict(list)

    def collector(name, doc):
        docs[name].append(doc)

    RE(plan(), collector)

    assert len(docs["event"]) == 3


def test_failing_describe_callback(RE, hw):
    class TestException(Exception):
        pass

    det = hw.det
    det2 = hw.det2

    def evil_cb(name, doc):
        if any(k in doc["data_keys"] for k in det.describe()):
            raise TestException

    RE.subscribe(evil_cb, "descriptor")

    def plan():
        yield Msg("open_run")
        try:
            yield Msg("declare_stream", None, det, name="det1")
            yield Msg("create", name="det1")
            yield Msg("read", det)
            yield Msg("save")
        finally:
            yield Msg("declare_stream", None, det2, name="det2")
            yield Msg("create", name="det2")
            yield Msg("read", det2)
            yield Msg("save")
        yield Msg("close_run")

    with pytest.raises(TestException):
        RE(plan())


def test_print_commands(RE):
    """test the printing of commands available.
    NOTE : An error here likely just means that this interface has changed.
    You'll likely want to change the test.
    (A test here is still a good idea as it at least raises awareness about
    the changes made breaking past API)
    """

    # testing the commands list
    commands1 = list(RE._command_registry.keys())
    commands2 = RE.commands

    assert commands1 == commands2

    # testing print commands
    # copy and paste most of the code...
    verbose = False
    print_command_reg1 = "List of available commands\n"
    for command, func in RE._command_registry.items():
        docstring = func.__doc__
        if verbose is False:
            docstring = docstring.split("\n")[0]
        print_command_reg1 += f"{command} : {docstring}\n"

    print_command_reg2 = RE.print_command_registry()
    assert print_command_reg1 == print_command_reg2


def test_broken_read_exception(RE):
    class Dummy:
        def __init__(self, name):
            self.name = name

        def read_configuration(self):
            return {}

        def describe_configuration(self):
            return {}

        def describe(self):
            return {}

        def read(self): ...

        def trigger(self): ...

    obj = Dummy("broken read")
    with pytest.raises(RuntimeError):
        RE([Msg("read", obj)])


def test_self_describe(RE):
    def inner():
        cls = yield Msg("RE_class")
        assert isinstance(RE, cls)

    RE(inner())


def test_thread_name(RE):
    "The RunEngine event loop should be on a thread with a given name."
    from ophyd.status import Status

    class MockDevice:
        name = "mock_device"

        def trigger(self):
            assert threading.current_thread().name == "bluesky-run-engine"
            return Status()

        def describe(self):
            return {}

        def read(self):
            return {}

    d = MockDevice()
    RE([Msg("trigger", d)])


def test_unsubscribe(RE):
    def foo(name, doc): ...

    for j in range(15):  # noqa: B007
        cid = RE.subscribe(foo)
        RE.unsubscribe(cid)

    assert len(RE.dispatcher._token_mapping) == 0


@pytest.fixture
def device_pair():
    from ophyd import StatusBase
    from ophyd.device import Device

    class MockDevice(Device):
        def stage(self):
            return StatusBase()

    return MockDevice(name="device1"), MockDevice(name="device2")


@pytest.mark.parametrize("set_finished", [True, False])
def test_wait_with_timeout(set_finished, RE, device_pair):
    def plan():
        status = yield Msg("stage", device_pair[0], group="test_group")
        if set_finished:
            status.set_finished()
        yield from wait(group="test_group", timeout=0.1)

    if set_finished:
        RE(plan())
    else:
        with pytest.raises(TimeoutError):
            RE(plan())


def test_wait_then_add_to_groups(RE, device_pair):
    def plan():
        status1 = yield Msg("stage", device_pair[0], group="test_group")
        done = yield from wait(group="test_group", timeout=0.001, error_on_timeout=False)
        assert done is False
        status2 = yield Msg("stage", device_pair[1], group="test_group")
        done = yield from wait(group="test_group", timeout=0.001, error_on_timeout=False)
        assert done is False
        status1.set_finished()
        status2.set_finished()
        done = yield from wait(group="test_group", timeout=0.001)
        assert done is True

    RE(plan())


async def finish_status_later(status, when: float, error: bool):
    await asyncio.sleep(when)
    if error:
        status.set_exception(ValueError("failed"))
    else:
        status.set_finished()


@pytest.mark.parametrize(
    "raises,exception,duration",
    [
        (True, FailedStatus, 0.1),
        (False, WaitForTimeoutError, 0.5),
    ],
)
def test_wait_when_one_status_returns_raises(RE, device_pair, raises, exception, duration):
    def plan():
        yield Msg("stage", device_pair[0], group="test_group")
        status2 = yield Msg("stage", device_pair[1], group="test_group")
        asyncio.create_task(finish_status_later(status2, 0.05, raises))
        yield from wait(group="test_group", timeout=0.5)

    start = ttime.monotonic()
    with pytest.raises(exception):
        RE(plan())
    elapsed = ttime.monotonic() - start
    assert elapsed == pytest.approx(duration, abs=0.1)


@pytest.mark.parametrize(
    "raises,exception,duration",
    [
        (True, FailedStatus, 0.1),
        (False, WaitForTimeoutError, 0.5),
    ],
)
def test_wait_returns_based_on_watch_raises(RE, device_pair, raises, exception, duration):
    def plan():
        yield Msg("stage", device_pair[0], group="test_group")
        stage2 = yield Msg("stage", device_pair[1], group="watch_group")
        asyncio.create_task(finish_status_later(stage2, 0.05, raises))
        yield from wait(group="test_group", watch=("watch_group",), timeout=0.5)

    start = ttime.monotonic()
    with pytest.raises(exception):
        RE(plan())
    elapsed = ttime.monotonic() - start
    assert elapsed == pytest.approx(duration, abs=0.1)


def test_set_finished_before_watch_completes_immediately(RE, device_pair):
    def plan():
        status = yield Msg("stage", device_pair[0], group="test_group")
        status.set_finished()
        yield Msg("stage", device_pair[1], group="watch_group")
        yield from wait(group="test_group", timeout=0.1, watch=("watch_group",))

    RE(plan())


def test_watch_finished_before_set_return_when_set_finishes(RE, device_pair):
    def plan():
        status1 = yield Msg("stage", device_pair[0], group="test_group")
        status2 = yield Msg("stage", device_pair[1], group="watch_group")
        asyncio.create_task(finish_status_later(status1, 0.2, False))
        asyncio.create_task(finish_status_later(status2, 0.05, False))
        yield from wait(group="test_group", timeout=0.3, watch=("watch_group",))

    start = ttime.monotonic()
    RE(plan())
    elapsed = ttime.monotonic() - start
    assert elapsed == pytest.approx(0.2, abs=0.05)


def test_set_failed_before_watch_raises_immediately(RE, device_pair):
    def plan():
        status = yield Msg("stage", device_pair[0], group="test_group")
        yield Msg("stage", device_pair[1], group="watch_group")
        status.set_exception(Exception("Mock device failed"))
        yield from wait(group="test_group", timeout=1.0, watch=("watch_group",))

    start = ttime.monotonic()
    with pytest.raises(FailedStatus):
        RE(plan())
    elapsed = ttime.monotonic() - start
    assert elapsed == pytest.approx(0.0, abs=0.1)


async def passing_coroutine():
    await asyncio.sleep(0.001)
    return 1080


async def failing_coroutine():
    await asyncio.sleep(0.001)
    raise RuntimeError("This is a runtime error")


def test_wait_with_planstub_coroutine(RE):
    def runtime_error_and_success_plan():
        tasks = yield from wait_for([lambda: passing_coroutine(), lambda: failing_coroutine()])
        passing_future, failing_future = tasks
        assert isinstance(passing_future, asyncio.Task) and isinstance(failing_future, asyncio.Task)
        assert passing_future.get_coro().__name__ == "passing_coroutine"
        assert failing_future.get_coro().__name__ == "failing_coroutine"
        assert passing_future.done() and passing_future.result() == 1080
        assert failing_future.done() and isinstance(failing_future.exception(), RuntimeError)

        with pytest.raises(RuntimeError, match="This is a runtime error"):
            raise failing_future.exception()  # the error is the same

    RE(runtime_error_and_success_plan())


def test_wait_with_planstub_status(RE):
    def runtime_error_and_success_plan():
        passing_task = asyncio.create_task(passing_coroutine())

        failing_task = asyncio.create_task(failing_coroutine())

        futures = yield from wait_for([lambda: passing_task, lambda: failing_task])
        assert passing_task is futures[0] and failing_task is futures[1]
        assert passing_task.done() and passing_task.result() == 1080
        assert failing_task.done() and isinstance(failing_task.exception(), RuntimeError)

    RE(runtime_error_and_success_plan())


def test_wait_with_planstub_timeout(RE):
    def timeout_plan():
        async def timeout_coro():
            await asyncio.sleep(1e9)

        timeout_task = asyncio.create_task(timeout_coro())
        try:
            yield from wait_for([lambda: timeout_task], timeout=0.01)
        finally:
            # Manually tear down the tasks, otherwise it can happen
            # after the test succeeds/fails
            timeout_task.cancel()

    with pytest.raises(WaitForTimeoutError, match="Plan failed to complete in the specified time"):
        RE(timeout_plan())


def test_1event_rewind(RE, hw):
    cache = []

    def hook(msg):
        cache.append(msg)

    doc_cache = []

    def cb(name, doc):
        doc_cache.append((name, doc))

    @run_decorator()
    def broken():
        yield from checkpoint()
        yield from trigger_and_read([hw.det], name="primary")
        yield from pause()
        yield from trigger_and_read([hw.det], name="primary")

    RE.msg_hook = hook
    with pytest.raises(RunEngineInterrupted):
        RE(broken(), cb)
    RE.resume()

    assert [_[1]["seq_num"] for _ in doc_cache if _[0] == "event"] == [1, 1, 2]

    assert [_[0] for _ in doc_cache] == [
        "start",
        "descriptor",
        "event",
        "event",
        "event",
        "stop",
    ]

    assert [msg.command for msg in cache] == [
        "open_run",
        "checkpoint",
        "trigger",
        "wait",
        "create",
        "read",
        "save",
        "pause",
        "trigger",
        "wait",
        "create",
        "read",
        "save",
        "trigger",
        "wait",
        "create",
        "read",
        "save",
        "close_run",
    ]


@requires_ophyd
def test_configure_multiple_descritpors(RE):
    from ophyd import Component as C
    from ophyd import Device, sim

    class SynWithConfig(Device):
        x = C(sim.Signal, value=0)
        y = C(sim.Signal, value=2)
        z = C(sim.Signal, value=3)

    det = SynWithConfig(name="det")
    det.x.name = "x"
    det.y.name = "y"
    det.z.name = "z"
    det.read_attrs = ["x"]
    det.configuration_attrs = ["y", "z"]

    @run_decorator()
    def plan(det):
        # run without an exciting descriptor
        yield from configure(det, {"z": 3})
        # force a descriptor to be created
        yield from declare_stream(det, name="primary")
        # and an event
        yield from trigger_and_read([det], name="primary")
        # update the config which will re-generate the descriptor
        yield from configure(det, {"z": 4})
        # and a second event
        yield from trigger_and_read([det], name="primary")

    d = DocCollector()
    RE(plan(det), d.insert)

    (start,) = d.start
    descA, descB = d.descriptor[start["uid"]]
    stop = d.stop[start["uid"]]

    assert descA["configuration"]["det"]["data"]["z"] == 3
    assert descB["configuration"]["det"]["data"]["z"] == 4

    for desc in (descA, descB):
        assert len(d.event[desc["uid"]]) == 1

    assert stop["num_events"]["primary"] == 2


def test_sync_scan_id_source(RE):
    def sync_scan_source(md: dict) -> int:
        return 314159

    RE.scan_id_source = sync_scan_source
    RE([Msg("open_run")])
    assert RE.md["scan_id"] == 314159


def test_async_scan_id_source(RE):
    async def async_scan_source(md: dict) -> int:
        return 42

    RE.scan_id_source = async_scan_source
    RE([Msg("open_run")])
    assert RE.md["scan_id"] == 42


@requires_ophyd
def test_descriptor_order(RE):
    from itertools import permutations

    from ophyd import Component, Device, Signal

    class Issue1930(Device):
        alpha = Component(Signal, value=1, kind="hinted")
        bravo = Component(Signal, value=2, kind="hinted")
        charlie = Component(Signal, value=3, kind="hinted")

    i1930 = Issue1930(name="i1930")

    for dets in permutations([i1930.alpha, i1930.bravo, i1930.charlie]):
        key_order = [d.name for d in dets]

        def check(key_order, name, doc):
            if name == "event":
                assert list(doc["data"]) == key_order
            elif name == "descriptor":
                assert list(doc["data_keys"]) == key_order

        RE(scan(dets, i1930.charlie, -1, 1, 2), lambda name, doc, key_order=key_order: check(key_order, name, doc))
