"""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, run_cache_key
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}


# -----------------------------------------------------------------------------
# The stage cache — a node whose inputs have not changed is not run again
# -----------------------------------------------------------------------------


class FakeCache:
    """A stage cache with no database behind it, and a record of what it was asked."""

    def __init__(self, entries: dict[str, dict | None] | None = None) -> None:
        self.entries = entries or {}
        self.asked: list[str] = []

    def lookup(self, key: str):
        self.asked.append(key)
        if key in self.entries:
            return True, self.entries[key]
        return False, None


def counting_node(flow: str = "study") -> tuple[Node, list[int]]:
    """A node that says how many times it actually ran."""
    calls: list[int] = []

    def train(lr, params):
        calls.append(1)
        return {"loss": lr * 2}

    node = make_node(
        "train", flow, train, requires=[spec("lr")], provides=[spec("loss")]
    )
    node.fingerprint = "fp-train"
    return node, calls


def test_a_cache_hit_restores_the_outputs_without_running_the_node():
    flow = double_flow()
    node, calls = counting_node()
    state = MemoryState()
    seen: list[NodeOutcome] = []
    key = run_cache_key("fp-train", {"study.lr": 0.5})
    cache = FakeCache({key: {"study.loss": 99.0}})

    Pipeline(nodes=[node], state=state, observer=seen.append, run_cache=cache).run(
        seed_values(flow, {"lr": 0.5})
    )

    assert calls == []
    # Restored into this run's own state, which is where everything
    # downstream of it looks — its namespace holds nothing otherwise.
    assert collect_result(flow, state) == {"loss": 99.0}
    assert seen[0].cached and seen[0].cache_key == key


def test_a_miss_runs_the_node_and_carries_what_would_be_stored():
    flow = double_flow()
    node, calls = counting_node()
    seen: list[NodeOutcome] = []
    cache = FakeCache()

    Pipeline(
        nodes=[node], state=MemoryState(), observer=seen.append, run_cache=cache
    ).run(seed_values(flow, {"lr": 0.5}))

    assert calls == [1]
    assert cache.asked == [run_cache_key("fp-train", {"study.lr": 0.5})]
    assert not seen[0].cached
    assert seen[0].cache_key and seen[0].output_values == {"study.loss": 1.0}


def test_the_key_follows_the_inputs():
    first = run_cache_key("fp", {"lr": 0.5})
    assert first != run_cache_key("fp", {"lr": 0.6})
    assert first != run_cache_key("other", {"lr": 0.5})
    assert first == run_cache_key("fp", {"lr": 0.5})


def test_an_artifact_input_counts_as_its_digest():
    digest = "sha256:" + "0" * 64
    # The same bytes under another name, of a size recorded differently, are
    # the same input — the reference is a handle, the digest is the content.
    assert run_cache_key(
        "fp", {"data": {"digest": digest, "name": "a.csv", "size": 3}}
    ) == run_cache_key("fp", {"data": {"digest": digest, "name": "b.csv", "size": 3}})


def test_a_node_with_no_fingerprint_is_never_looked_up():
    """Built-in and connector nodes, and anything declared `cache=False`."""
    flow = double_flow()
    node, calls = counting_node()
    node.fingerprint = ""
    seen: list[NodeOutcome] = []
    cache = FakeCache()

    Pipeline(
        nodes=[node], state=MemoryState(), observer=seen.append, run_cache=cache
    ).run(seed_values(flow, {"lr": 0.5}))

    assert calls == [1]
    assert cache.asked == []
    assert seen[0].cache_key == ""