"""What the decorators declare, and what the engine would be given."""

import pytest

from fluksio.flow.schemas import FlowDef
from fluksio.sdk import MARKER, Flow, Port, SyncError, node, use

# The functions live here rather than in each test: a node's module is what the
# generated body imports, and `__main__` is refused at sync for that reason.


@node(provides=[Port("dataset", "artifact"), Port("rows", "int")])
def prepare(limit=8):
    return {"dataset": {"digest": "sha256:x", "size": 1}, "rows": limit}


@node(
    requires=["dataset", Port("lr", "float")],
    provides=[Port("loss", "float", stream=True), Port("weights", "artifact")],
    device="gpu",
    device_policy="prefer",
)
def fit(dataset, lr, epochs=3):
    for step in range(epochs):
        yield {"loss": 1.0 / (step + 1)}
    return {"weights": {"digest": "sha256:y", "size": 1}}


@node(requires=["weights"], provides=Port("score", "float"))
def evaluate(weights):
    return 0.5


@node(requires=["dataset"], provides=[Port("augmented", "artifact")])
def augment(dataset):
    return {"augmented": dataset}


def a_flow(name="train", **kwargs):
    options = {
        "nodes": [prepare, fit, evaluate],
        "inputs": [Port("lr", "float", initial=0.01)],
        "outputs": ["score"],
    }
    options.update(kwargs)
    return Flow(name, **options)


def test_document_is_a_flow_the_engine_accepts():
    document = a_flow().document({"kind": "python", "repo": "/r", "commit": "abc"})
    definition = FlowDef.model_validate(document)

    assert definition.mode == "batch"
    assert [n.id for n in definition.nodes] == ["prepare", "fit", "evaluate"]
    assert definition.origin is not None
    assert definition.origin.commit == "abc"
    assert [i.spec.name for i in definition.inputs] == ["lr"]
    assert definition.outputs == ["score"]


def test_defaults_become_settings_and_ports_do_not():
    definition = FlowDef.model_validate(a_flow().document())
    nodes = {n.id: n for n in definition.nodes}

    assert nodes["prepare"].params == {"limit": 8}
    # `dataset` and `lr` are ports, so they are not settings as well.
    assert nodes["fit"].params == {"epochs": 3}
    assert {p.port for p in nodes["fit"].requires} == {"dataset", "lr"}


def test_device_travels_to_the_node():
    definition = FlowDef.model_validate(a_flow().document())
    node_def = next(n for n in definition.nodes if n.id == "fit")

    assert node_def.device == "gpu"
    assert node_def.device_policy == "prefer"


def test_a_required_port_takes_the_type_it_is_provided_as():
    """`requires=["dataset"]` is not a second declaration of its type."""
    definition = FlowDef.model_validate(a_flow().document())
    node_def = next(n for n in definition.nodes if n.id == "fit")

    assert next(p for p in node_def.requires if p.port == "dataset").dtype == "artifact"


def test_use_rewires_and_reconfigures_one_flow_only():
    train = a_flow()
    finetune = Flow(
        "finetune",
        nodes=[prepare, augment, use(fit, wire={"dataset": "augmented"}, epochs=9)],
        inputs=[Port("lr", "float")],
    )

    theirs = next(n for n in finetune.document()["nodes"] if n["id"] == "fit")
    ours = next(n for n in train.document()["nodes"] if n["id"] == "fit")
    assert [p["name"] for p in theirs["requires"]] == ["augmented", "lr"]
    assert theirs["params"] == {"epochs": 9}
    # The same function in another flow is untouched by that.
    assert [p["name"] for p in ours["requires"]] == ["dataset", "lr"]
    assert ours["params"] == {"epochs": 3}


def unknown_port(a):
    return a


def two_arguments(a, b):
    return a


def one_default(a=1):
    return a


def a_generator():
    yield {"x": 1}


def reads_dataset(dataset):
    return dataset


def test_a_port_with_no_parameter_is_refused():
    with pytest.raises(SyncError, match="no parameter to arrive in"):
        node(requires=["nope"])(unknown_port)


def test_a_parameter_that_is_neither_port_nor_setting_is_refused():
    with pytest.raises(SyncError, match="neither a port nor a setting"):
        node(requires=["a"])(two_arguments)


def test_a_name_the_store_would_refuse_is_refused_here():
    with pytest.raises(SyncError, match="node id"):
        node(requires=["a"], id="Trainer")(one_default)


def test_a_generator_cannot_use_the_bare_return_shorthand():
    with pytest.raises(SyncError, match="yields"):
        node(provides=Port("x", "float"))(a_generator)


def test_a_setting_that_is_also_a_port_is_refused():
    with pytest.raises(SyncError, match="both a port and a setting"):
        Flow("clash", nodes=[use(prepare, limit=2), _sets("dataset")])


def test_a_setting_that_is_not_a_parameter_is_refused():
    with pytest.raises(SyncError, match="not a parameter"):
        Flow("stray", nodes=[use(prepare, nonesuch=1)])


def _sets(name):
    """A node whose settings clash with its own port, built for the test."""
    fn = node(requires=["dataset"], settings={name: 1})
    return fn(reads_dataset)


def test_a_message_nobody_provides_is_refused():
    orphan = node(requires=["missing"], id="orphan")(reads_dataset_missing)
    with pytest.raises(SyncError, match="which no node in it provides"):
        Flow("broken", nodes=[orphan])


def reads_dataset_missing(missing):
    return missing


def test_an_output_nobody_produces_is_refused():
    with pytest.raises(SyncError, match="not produced by any"):
        a_flow("typo", outputs=["scoer"])


def test_two_nodes_of_one_name_are_refused():
    with pytest.raises(SyncError, match="tell them apart"):
        Flow("twice", nodes=[prepare, prepare])


def test_a_type_both_sides_disagree_on_is_refused():
    wrong = node(requires=[Port("dataset", "float")], id="wrong")(reads_dataset)
    with pytest.raises(SyncError, match="provided as artifact"):
        Flow("mistyped", nodes=[prepare, wrong])


def test_shim_imports_rather_than_copies():
    shims = a_flow().shims()

    assert shims["prepare"].startswith(MARKER)
    assert "from tests.sdk.test_build import prepare as _impl" in shims["prepare"]
    assert "def process(**settings):" in shims["prepare"]
    assert "return _impl(**settings)" in shims["prepare"]


def test_shim_of_a_generator_delegates_and_keeps_its_return_value():
    """A bare `yield from` streams but drops what the generator returns."""
    assert (
        "return (yield from _impl(dataset=dataset, lr=lr, **settings))"
        in a_flow().shims()["fit"]
    )


def test_a_generator_shim_publishes_both_the_stream_and_the_result():
    """Driven the way the worker drives it: every yield, then the return."""
    namespace: dict = {}
    exec(compile(a_flow().shims()["fit"], "<shim>", "exec"), namespace)

    generator = namespace["process"](dataset=None, lr=0.5, epochs=2)
    streamed = []
    result = None
    while True:
        try:
            streamed.append(next(generator))
        except StopIteration as stop:
            result = stop.value
            break

    assert [step["loss"] for step in streamed] == [1.0, 0.5]
    # A bare `yield from` would stream the same and lose this.
    assert result == {"weights": {"digest": "sha256:y", "size": 1}}


def test_shim_of_a_single_port_wraps_the_bare_return():
    assert "return {'score': _impl(weights=weights" in a_flow().shims()["evaluate"]


@node(requires=["scale"], provides=Port("scaled", "float"))
def scale(scale):
    return scale * 2.0


def test_a_port_named_like_its_function_still_calls_the_function():
    """The body imports under an alias, so the port cannot shadow the function."""
    flow = Flow("collide", nodes=[scale], inputs=[Port("scale", "float", initial=1.0)])
    namespace: dict = {}
    exec(compile(flow.shims()["scale"], "<shim>", "exec"), namespace)

    assert namespace["process"](scale=2.0) == {"scaled": 4.0}


def reads_settings(settings):
    return settings


def test_a_port_the_generated_body_reserves_is_refused():
    refused = node(requires=["settings"], provides=Port("out", "float"))(reads_settings)
    with pytest.raises(SyncError, match="keeps for itself"):
        Flow("reserved", nodes=[refused], inputs=[Port("settings", "float")])


def test_every_shim_compiles_and_defines_process():
    for code in a_flow().shims().values():
        namespace: dict = {}
        exec(compile(code, "<shim>", "exec"), namespace)
        assert callable(namespace["process"])


def test_the_decorators_leave_the_function_alone():
    """The point of the whole feature: it is still your code."""
    assert prepare(limit=2)["rows"] == 2
    assert [step["loss"] for step in fit(None, 0.5, epochs=2)] == [1.0, 0.5]
    assert evaluate(None) == 0.5


# -----------------------------------------------------------------------------
# What a function emits against what it declares
#
# The engine refuses an undeclared key at the first yield, which is right but
# late — a sweep can be an hour in. A literal one is a typo, and a typo is
# readable from the source.
# -----------------------------------------------------------------------------


def trains_with_a_typo(steps=3):
    for _ in range(steps):
        yield {"lss": 0.5}
    return {"final_loss": 0.5}


def trains(steps=3):
    for _ in range(steps):
        yield {"loss": 0.5}
    return {"final_loss": 0.5}


def emits_a_typo(steps=3):
    import fluksio

    fluksio.emit(lss=0.5)
    return {"final_loss": 0.5}


def yields_a_name_it_computes(steps=3):
    for index in range(steps):
        yield {f"loss_{index}": 0.5}
    return {"final_loss": 0.5}


def trains_beside_a_helper(steps=3):
    def every_pair():
        yield {"internal": 1}

    for _ in range(steps):
        yield {"loss": 0.5}
    return {"final_loss": 0.5}


def test_a_yielded_key_no_port_declares_is_refused():
    with pytest.raises(SyncError, match="lss"):
        node(provides=[Port("loss", "float"), Port("final_loss", "float")])(
            trains_with_a_typo
        )


def test_an_emitted_key_no_port_declares_is_refused():
    with pytest.raises(SyncError, match="lss"):
        node(provides=[Port("final_loss", "float")])(emits_a_typo)


def test_declared_keys_pass():
    assert node(provides=[Port("loss", "float"), Port("final_loss", "float")])(trains)


def test_a_key_the_code_computes_is_left_to_the_engine():
    """Only literals are readable here; the rest is checked where it runs."""
    assert node(provides=[Port("final_loss", "float")])(yields_a_name_it_computes)


def test_a_helper_defined_inside_the_node_is_not_the_nodes_ports():
    assert node(provides=[Port("loss", "float"), Port("final_loss", "float")])(
        trains_beside_a_helper
    )


def test_a_negative_timeout_is_refused():
    with pytest.raises(SyncError, match="0 or more"):
        node(requires=["a"], timeout=-1)(one_default)


def test_a_zero_timeout_means_no_limit():
    decorated = node(requires=["a"], timeout=0)(one_default)
    assert decorated.__fluksio__.timeout == 0


def test_a_node_in_a_script_run_directly_is_refused_at_sync_not_at_import():
    """A study module has to be runnable as a script for a self-check.

    The refusal belongs where the body is generated: the decorator hands the
    function back untouched, so `python study.py` declares its nodes, calls
    them and checks itself. What cannot be done is syncing them, because
    nothing could import `__main__`.
    """

    def probe(rows=1):
        return {"score": float(rows)}

    probe.__module__ = "__main__"
    decorated = node(provides=[Port("score", "float")])(probe)
    # Declared and callable, which is the whole point of running the file.
    assert decorated(rows=2) == {"score": 2.0}

    flow = Flow("mainflow", nodes=[decorated], outputs=["score"])
    with pytest.raises(SyncError, match="run directly"):
        flow.document()
    with pytest.raises(SyncError, match="run directly"):
        flow.shims()