"""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 app.flow.messages import DType, MessageSpec
from app.flow.nodes import Node
from app.flow.pipeline import Pipeline
from app.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_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