"""Deadlines retain worker capacity when ordinary work is blocked."""

import threading
import time

import pytest

from fluksio.flow import executor
from fluksio.flow.executor import ExecutionService
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.pipeline import Pipeline
from fluksio.flow.queue import MemoryWorkQueue, WorkItem


def _node(name, function, requires=(), provides=()):
    node = Node(
        f=function,
        requires=[MessageSpec(name=key, dtype=DType.FLOAT) for key in requires],
        provides=[MessageSpec(name=key, dtype=DType.FLOAT) for key in provides],
        name=name,
    )
    node.assign_flow("f", name)
    return node


@pytest.mark.parametrize("reclaimed", [False, True])
def test_a_due_node_runs_while_ordinary_nodes_and_cascades_are_full(
    monkeypatch, reclaimed
):
    started, release, fired, recovered, touched = (threading.Event() for _ in range(5))
    calls = []

    class Queue(MemoryWorkQueue):
        def reclaim_stale(self, min_idle_ms):
            if not reclaimed or not started.is_set() or recovered.is_set():
                return []
            recovered.set()
            return [
                WorkItem(
                    kind="cascade",
                    node="f.source",
                    flow="f",
                    outputs={"f.raw": float(i)},
                    entry_id=f"recovered-{i}",
                )
                for i in range(3)
            ]

        def touch(self, entry_ids):
            if "recovered-0" in entry_ids:
                touched.set()

    def slow(raw, params):
        calls.append(raw)
        started.set()
        assert release.wait(5)

    source = _node("source", lambda params: None, provides=["raw"])
    slow_node = _node("slow", slow, requires=["raw"])
    timer = _node("timer", lambda params: None, provides=["stop"])
    stop = _node("stop", lambda stop, params: fired.set(), requires=["stop"])
    queue = Queue()
    service = ExecutionService(queue, max_workers=1, max_cascades=1)
    pipeline = Pipeline(
        [source, slow_node, timer, stop], work_queue=queue, node_pool=service.node_pool
    )
    service.bind(pipeline)
    monkeypatch.setattr(executor, "DELAYED_INTERVAL_S", 0.02)
    monkeypatch.setattr(executor, "RECLAIM_INTERVAL_S", 0.02)
    monkeypatch.setattr(executor, "TOUCH_INTERVAL_S", 0.02)
    service.start()
    try:
        source.inject({"raw": 1.0})
        assert started.wait(1)
        if reclaimed:
            assert touched.wait(1), "reclaimed work blocked the timer's heartbeat"
        else:
            for i in range(3):
                source.inject({"raw": float(i)})
        queue.add_delayed(
            WorkItem(kind="cascade", node=timer.id, flow="f", outputs={"f.stop": 1.0}),
            time.time(),
        )
        assert fired.wait(2), "the due node waited for an ordinary worker"
        assert len(calls) == 1, "ordinary work used the timer reserve"
    finally:
        release.set()
        service.stop()


def test_timers_can_use_all_idle_normal_workers():
    """A reserve must not cap an otherwise idle engine at two workers."""
    service = ExecutionService(MemoryWorkQueue(), max_workers=3, max_cascades=3)
    started = [threading.Event() for _ in range(3)]
    release = threading.Event()

    def block(index):
        started[index].set()
        assert release.wait(3)

    futures = []
    try:
        for i in range(3):
            futures.append(service._node_pools.due.submit(block, i))
        assert all(event.wait(1) for event in started)
        assert service._node_pools.available() == (0, executor.DUE_RESERVE)
    finally:
        release.set()
        for future in futures:
            future.result(1)
        service.stop()