"""What a run is made of: isolation, the per-node record, and what it reports.

The service itself needs a database, so what is checked here is the part that
decides whether a run is correct — that two runs of one flow cannot see each
other's messages, that every node executed is reported once, and that the
parameters a caller sends are refused before anything runs if they are wrong.
"""

import pytest

from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.pipeline import NodeOutcome, Pipeline
from fluksio.flow.runs import (
    MetricSink,
    RunRejected,
    batch_issues,
    collect_result,
    digest_of,
    required_labels,
    seed_values,
)
from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef
from fluksio.flow.state import MemoryState


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


def make_node(node_id: str, flow: str, f, requires=(), provides=()) -> Node:
    node = Node(f=f, requires=list(requires), provides=list(provides), name=node_id)
    node.assign_flow(flow, node_id)
    return node


def double_flow() -> FlowDef:
    """A flow with one input, one node and one declared output."""
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=spec("lr"), initial=0.1)],
        outputs=["loss"],
        nodes=[
            NodeDef(
                id="train",
                requires=[spec("lr")],
                provides=[spec("loss")],
            )
        ],
    )


def build(flow: FlowDef, state: MemoryState, observer=None) -> Pipeline:
    """The pipeline a run drives, without the controller that normally builds it."""
    node = make_node(
        "train",
        flow.name,
        lambda lr, params: {"loss": lr * 2},
        requires=[spec("lr")],
        provides=[spec("loss")],
    )
    return Pipeline(nodes=[node], state=state, observer=observer)


# -----------------------------------------------------------------------------
# Isolation — the reason a run has a state backend of its own
# -----------------------------------------------------------------------------


def test_two_runs_of_one_flow_do_not_see_each_other():
    flow = double_flow()
    first, second = MemoryState(), MemoryState()

    build(flow, first).run(seed_values(flow, {"lr": 0.5}))
    build(flow, second).run(seed_values(flow, {"lr": 4.0}))

    assert collect_result(flow, first) == {"loss": 1.0}
    assert collect_result(flow, second) == {"loss": 8.0}


def test_result_falls_back_to_everything_the_flow_holds():
    flow = double_flow()
    flow.outputs = []
    state = MemoryState()

    build(flow, state).run(seed_values(flow, {"lr": 1.0}))

    # The input is part of what the flow ended up holding; the engine's own
    # bookkeeping keys are not.
    assert collect_result(flow, state) == {"lr": 1.0, "loss": 2.0}


# -----------------------------------------------------------------------------
# The record — every node a run executed, reported once
# -----------------------------------------------------------------------------


def test_observer_sees_every_node_once():
    flow = double_flow()
    seen = []

    build(flow, MemoryState(), observer=seen.append).run(seed_values(flow, {"lr": 1.0}))

    assert [(o.node, o.ok) for o in seen] == [("study.train", True)]
    assert seen[0].outputs == 1


def test_observer_reports_a_failing_node_with_its_error():
    seen = []

    def boom(params):
        raise ValueError("no convergence")

    node = make_node("train", "study", boom, provides=[spec("loss")])
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    assert len(seen) == 1
    assert not seen[0].ok
    assert "no convergence" in seen[0].error


def test_a_failing_observer_does_not_take_the_node_down():
    def refuse(_outcome):
        raise RuntimeError("the database is gone")

    node = make_node(
        "train", "study", lambda params: {"loss": 1.0}, provides=[spec("loss")]
    )
    pipeline = Pipeline(nodes=[node], state=MemoryState(), observer=refuse)
    pipeline.run()

    assert pipeline.state["study.loss"] == 1.0


# -----------------------------------------------------------------------------
# Producing values before returning
#
# A number worth keeping is an output, not a log. A node that produces over
# time is a generator, and each yield is published on the port it names — so
# the run's metrics are its streaming outputs rather than something recorded
# beside them.
# -----------------------------------------------------------------------------


def training_flow() -> FlowDef:
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=spec("steps", DType.INT), initial=3)],
        outputs=["final_loss"],
        nodes=[
            NodeDef(
                id="train",
                requires=[spec("steps", DType.INT)],
                provides=[
                    spec("loss", stream=True),
                    spec("final_loss"),
                ],
            )
        ],
    )


def training_node(flow: str = "study") -> Node:
    def train(steps, params):
        loss = 1.0
        for _ in range(steps):
            loss = loss / 2
            yield {"loss": loss}
        return {"final_loss": loss}

    return make_node(
        "train",
        flow,
        train,
        requires=[spec("steps", DType.INT)],
        provides=[spec("loss", stream=True), spec("final_loss")],
    )


def test_each_yield_is_published_and_the_return_is_the_result():
    flow = training_flow()
    state = MemoryState()
    seen: list[tuple[str, dict]] = []
    pipeline = Pipeline(
        nodes=[training_node()],
        state=state,
        emission_observer=lambda node, outputs: seen.append((node, outputs)),
    )

    pipeline.run(seed_values(flow, {"steps": 3}))

    # It returned something, so every yield was a value produced on the way —
    # each published on `loss` the moment it happened.
    assert [outputs["study.loss"] for _node, outputs in seen] == [0.5, 0.25, 0.125]
    assert all(node == "study.train" for node, _ in seen)
    # The latest of them is what the message holds, as for any producer.
    assert state["study.loss"] == 0.125
    # And what it returned is the node's output, and so the run's result.
    assert collect_result(flow, state) == {"final_loss": 0.125}


def test_without_a_return_the_last_yield_is_the_result():
    def train(params):
        yield {"loss": 1.0}
        yield {"loss": 0.5}

    node = make_node("train", "study", train, provides=[spec("loss", stream=True)])
    seen: list[tuple[str, dict]] = []
    pipeline = Pipeline(
        nodes=[node],
        state=MemoryState(),
        emission_observer=lambda n, o: seen.append((n, o)),
    )
    pipeline.run()

    # The last yield is the node's output rather than an emission, so it is
    # not counted twice.
    assert [outputs["study.loss"] for _node, outputs in seen] == [1.0]
    assert pipeline.state["study.loss"] == 0.5


def test_emissions_are_checked_against_the_port_they_name():
    def wrong(params):
        yield {"loss": "not a number"}
        return {"final_loss": 1.0}

    node = make_node(
        "train",
        "study",
        wrong,
        provides=[spec("loss", stream=True), spec("final_loss")],
    )
    seen: list[NodeOutcome] = []
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    # A wrong type is a failed node, exactly as it is for a return value —
    # which is the point of emissions going out through declared ports.
    assert not seen[0].ok
    assert "loss" in seen[0].error


def test_an_emission_that_nothing_declares_is_ignored():
    def stray(params):
        yield {"undeclared": 1.0}
        return {"final_loss": 2.0}

    node = make_node("train", "study", stray, provides=[spec("final_loss")])
    state = MemoryState()
    Pipeline(nodes=[node], state=state).run()

    assert "study.undeclared" not in state
    assert state["study.final_loss"] == 2.0


def test_emissions_reach_the_run_as_a_series_with_a_step_each():
    sink = MetricSink("run-1", batch=1)
    sink.handle("study.train", {"study.loss": 1.0, "study.tag": "ignored"})
    sink.handle("study.train", {"study.loss": 0.5})

    # Numbers become the run's series; anything else is on the run some other
    # way — as its result, or as an artifact.
    assert sink._steps == {"study.loss": 1}


# -----------------------------------------------------------------------------
# What a caller may ask for
# -----------------------------------------------------------------------------


def test_parameters_must_be_declared_inputs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="not an input"):
        seed_values(flow, {"learning_rate": 0.1})


def test_parameters_are_type_checked_before_anything_runs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="lr"):
        seed_values(flow, {"lr": "fast"})


def test_a_rate_limited_port_cannot_be_run_as_a_batch():
    flow = double_flow()
    flow.nodes[0].provides = [spec("loss", interval=30)]
    # Without a queue there is no timer to release what an interval holds, so
    # the value would be dropped rather than delayed.
    assert batch_issues(flow)
    assert not batch_issues(double_flow())


def test_a_streaming_port_may_thin_itself_out():
    flow = double_flow()
    flow.nodes[0].provides = [spec("loss", interval=0.5, stream=True)]
    # On a curve, an interval is asking for the canvas not to be flooded —
    # the run's history still keeps every value.
    assert not batch_issues(flow)


def test_the_digest_identifies_the_inputs_not_their_order():
    assert digest_of({"a": 1, "b": 2}, 3) == digest_of({"b": 2, "a": 1}, 3)
    assert digest_of({"a": 1}, 3) != digest_of({"a": 1}, 4)


def test_labels_come_from_the_nodes_that_ask_for_a_device():
    flow = double_flow()
    assert required_labels(flow) == []
    flow.nodes[0].device = "gpu"
    assert required_labels(flow) == ["gpu"]


def test_a_preferred_device_does_not_hold_a_run_back():
    flow = double_flow()
    flow.nodes[0].device = "gpu"
    flow.nodes[0].device_policy = "prefer"
    # It runs on the engine when no such worker is attached, so waiting for
    # one would be waiting for something the run does not need.
    assert required_labels(flow) == []


def test_a_runs_seed_fills_an_input_of_that_name():
    flow = double_flow()
    flow.inputs.append(FlowInput(spec=spec("seed", DType.INT), initial=0))

    # Otherwise the field that tells two runs of one configuration apart would
    # only look like the number the flow draws from.
    assert seed_values(flow, {}, seed=7)["study.seed"] == 7
    # An explicit parameter still wins.
    assert seed_values(flow, {"seed": 3}, seed=7)["study.seed"] == 3


def test_a_flow_without_a_seed_input_ignores_the_runs_seed():
    flow = double_flow()
    assert seed_values(flow, {"lr": 1.0}, seed=7) == {"study.lr": 1.0}