"""Synchronous nodes wait until every input is fresh.

The ordering guarantee at the stateful boundary rests on the state backend's
atomic operations, so those are covered here too.
"""

from concurrent.futures import ThreadPoolExecutor

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


def spec(name: str) -> MessageSpec:
    return MessageSpec(name=name, dtype=DType.FLOAT)


def node(node_id: str, f, requires=(), provides=(), params=None) -> Node:
    n = Node(
        f=f,
        requires=list(requires),
        provides=list(provides),
        params=params or {},
        name=node_id,
    )
    n.assign_flow("f", node_id)
    return n


def test_synchronous_node_waits_for_all_inputs_to_be_fresh():
    runs: list[str] = []

    def sensor_a(params):
        return {"a": 1.0}

    def sensor_b(params):
        return {"b": 2.0}

    def sync(a, b, params):
        runs.append("sync")
        return None

    def eager(a, b, params):
        runs.append("eager")
        return None

    a = node("a", sensor_a, provides=[spec("a")])
    b = node("b", sensor_b, provides=[spec("b")])
    sync_node = node(
        "sync", sync, requires=[spec("a"), spec("b")], params={"synchronous": True}
    )
    eager_node = node("eager", eager, requires=[spec("a"), spec("b")])

    Pipeline(nodes=[a, b, sync_node, eager_node], max_workers=1)

    a.inject()
    assert runs == []  # b has never arrived

    b.inject()
    assert runs.count("eager") == 1
    assert runs.count("sync") == 1

    # Only a is new: the eager node runs again, the synchronous one waits.
    runs.clear()
    a.inject()
    assert runs == ["eager"]

    # Now b is new as well, so both have moved on.
    runs.clear()
    b.inject()
    assert sorted(runs) == ["eager", "sync"]


def test_a_seeded_value_counts_as_having_arrived():
    """A re-put flow used to come back with its synchronous nodes wedged.

    Seeding wrote the value and not its version, and version 0 reads as
    "never published" — so the node reported active and ok and never ran
    again, with the value it was waiting for sitting right there in state.
    """
    runs: list[str] = []

    def sync(a, b, params):
        runs.append("sync")
        return None

    def sensor_b(params):
        return {"b": 2.0}

    b = node("b", sensor_b, provides=[spec("b")])
    sync_node = node(
        "sync", sync, requires=[spec("a"), spec("b")], params={"synchronous": True}
    )

    Pipeline(
        nodes=[b, sync_node],
        max_workers=1,
        initial_values={"f.a": 1.0},
    )

    b.inject()
    assert runs == ["sync"]


def test_increment_and_multi_get():
    state = MemoryState()

    assert state.increment("counter") == 1
    assert state.increment("counter") == 2

    state.update({"a": 1, "b": 2})
    assert state.get_multi(["a", "b", "missing"]) == {
        "a": 1,
        "b": 2,
        "missing": None,
    }


def test_compare_and_swap_only_applies_on_match():
    state = MemoryState()
    state.update({"a": 1, "b": 2})

    assert state.compare_and_swap_multi({"a": 1, "b": 2}, {"a": 10, "new": 100})
    assert state.get("a") == 10
    assert state.get("new") == 100

    assert not state.compare_and_swap_multi({"a": 1}, {"a": 99})
    assert state.get("a") == 10


def test_only_one_thread_wins_a_compare_and_swap():
    state = MemoryState()
    state.set("version", 1)

    def claim() -> bool:
        return state.compare_and_swap_multi({"version": 1}, {"version": 2})

    with ThreadPoolExecutor(max_workers=8) as pool:
        results = list(pool.map(lambda _: claim(), range(8)))

    assert sum(results) == 1