"""Scheduling work scales with the graph and observes runtime gates."""

import threading
from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext

import pytest

from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.pipeline import Pipeline


def _node(name, function, requires=(), provides=(), flow="f"):
    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(flow, name)
    return node


def test_a_long_chain_checks_each_candidate_once():
    class CountingPipeline(Pipeline):
        checks = 0

        def _gate_blocks(self, node):
            self.checks += 1
            return super()._gate_blocks(node)

    nodes = []
    for i in range(100):

        def relay(params, _index=i, **inputs):
            return {f"v{_index + 1}": inputs[f"v{_index}"] + 1.0}

        nodes.append(_node(f"n{i}", relay, requires=[f"v{i}"], provides=[f"v{i + 1}"]))
    pipeline = CountingPipeline(nodes)

    pipeline.run({"f.v0": 0.0})

    assert pipeline.values()["f.v100"]["value"] == 100.0
    assert pipeline.checks == len(nodes)


@pytest.mark.parametrize("disabled", [False, True])
def test_a_gate_opened_mid_wave_releases_its_waiting_candidate(disabled):
    running, release, held, acted = (threading.Event() for _ in range(4))

    class WatchingPipeline(Pipeline):
        def _gate_blocks(self, node):
            blocked = super()._gate_blocks(node)
            if blocked:
                held.set()
            return blocked

    def slow(params):
        running.set()
        assert release.wait(3)

    pipeline = WatchingPipeline(
        [
            _node("slow", slow, flow="running"),
            _node("action", lambda params: acted.set(), flow="paused"),
        ]
    )
    if disabled:
        pipeline.set_disabled({"paused"})
    else:
        pipeline.pause("paused")
    with ThreadPoolExecutor(max_workers=1) as driver:
        future = driver.submit(pipeline.run)
        try:
            assert running.wait(1) and held.wait(1)
            if disabled:
                pipeline.set_disabled(set())
            with nullcontext() if disabled else pipeline.stepping("paused"):
                release.set()
                future.result(2)
            assert acted.is_set()
        finally:
            release.set()