"""The stage cache, from the side that needs a database.

The pipeline half — what a hit restores and what a key is made of — is in
`tests/flow/test_runs.py`, which runs without one.
"""

import json
from datetime import UTC, datetime

import pytest
from sqlmodel import Session

from fluksio.core.db import engine as db_engine
from fluksio.flow.artifacts import ArtifactStore
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.pipeline import NodeOutcome
from fluksio.flow.runs import (
    OUTPUT_CAP,
    RunCache,
    RunRejected,
    _cacheable,
    new_run_id,
    resolve_references,
    seed_values,
)
from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef
from fluksio.models import Run, RunArtifact, RunNode


def test_a_run_cache_finds_what_an_earlier_run_recorded(tmp_path):
    """The run history is the cache; there is no second store to keep."""
    store = ArtifactStore(tmp_path / "artifacts")
    reference = store.put([b"payload"], name="data.bin")
    plain, with_artifact, collected = "k-plain", "k-artifact", "k-collected"
    with Session(db_engine) as session:
        session.add(
            RunNode(
                run_id="cache-1",
                node="study.a",
                status="ok",
                cache_key=plain,
                outputs=json.dumps({"study.loss": 1.5}),
            )
        )
        session.add(
            RunNode(
                run_id="cache-2",
                node="study.b",
                status="ok",
                cache_key=with_artifact,
                outputs=json.dumps({"study.data": reference}),
            )
        )
        session.add(
            RunNode(
                run_id="cache-3",
                node="study.c",
                status="ok",
                cache_key=collected,
                outputs=json.dumps(
                    {"study.data": {**reference, "digest": "sha256:" + "1" * 64}}
                ),
            )
        )
        session.commit()

    cache = RunCache(store)
    assert cache.lookup(plain) == (True, {"study.loss": 1.5})
    assert cache.lookup(with_artifact) == (True, {"study.data": reference})
    # Its bytes have gone from the store, so the reference names nothing a
    # restored run could open. That is a miss, not a broken run.
    assert cache.lookup(collected) == (False, None)
    assert cache.lookup("never-seen") == (False, None)
    assert cache.lookup("") == (False, None)


def test_what_may_be_stored_as_a_cache_entry():
    """A row carries a key and its outputs together, or neither."""
    ok = NodeOutcome(
        node="study.a", ok=True, cache_key="k", output_values={"study.loss": 1.0}
    )
    assert _cacheable(ok) == '{"study.loss":1.0}'
    # A node that published nothing is still an answer worth reusing.
    assert _cacheable(ok.model_copy(update={"output_values": None})) == "null"
    # Not cacheable: it failed, it has no key, or it returned too much.
    assert _cacheable(ok.model_copy(update={"ok": False})) is None
    assert _cacheable(ok.model_copy(update={"cache_key": ""})) is None
    big = {"study.data": "x" * (OUTPUT_CAP + 1)}
    assert _cacheable(ok.model_copy(update={"output_values": big})) is None


# -----------------------------------------------------------------------------
# Naming an artifact from outside the process that made it
#
# A python caller passes the reference it holds. A shell holds nothing, so the
# same input also takes `@run:<id>.<output>` or a bare digest, resolved here
# rather than in each client.
# -----------------------------------------------------------------------------


def artifact_flow() -> FlowDef:
    """A flow taking a dataset somebody else's run produced."""
    dataset = MessageSpec(name="dataset", dtype=DType.ARTIFACT)
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=dataset)],
        nodes=[NodeDef(id="train", requires=[dataset])],
    )


@pytest.fixture
def made_artifact():
    """A finished run with one artifact, as a later run would find it."""
    digest = "sha256:" + "a1" * 32
    reference = {
        "digest": digest,
        "size": 12,
        "media_type": "text/csv",
        "name": "cities.csv",
    }
    run_id = new_run_id()
    with Session(db_engine) as session:
        session.add(
            Run(
                id=run_id,
                flow="prepare",
                status="ok",
                result={"dataset": reference},
                created_at=datetime.now(UTC),
            )
        )
        session.add(
            RunArtifact(
                run_id=run_id,
                name="prepare.dataset",
                node="load",
                digest=digest,
                size=12,
            )
        )
        session.commit()
    yield run_id, reference
    with Session(db_engine) as session:
        session.delete(session.get(RunArtifact, (run_id, "prepare.dataset")))
        session.delete(session.get(Run, run_id))
        session.commit()


def test_a_run_reference_resolves_to_what_that_run_produced(made_artifact):
    run_id, reference = made_artifact
    resolved = resolve_references(artifact_flow(), {"dataset": f"@run:{run_id}.dataset"})

    # The producer's own reference, file name and all — not one rebuilt from
    # the row, which carries the message name instead.
    assert resolved["dataset"] == reference


def test_a_bare_digest_resolves_to_the_bytes_under_it(made_artifact):
    _run_id, reference = made_artifact
    resolved = resolve_references(artifact_flow(), {"dataset": reference["digest"]})

    assert resolved["dataset"]["digest"] == reference["digest"]
    assert resolved["dataset"]["size"] == 12


def test_a_resolved_reference_passes_the_input_check(made_artifact):
    run_id, _reference = made_artifact
    flow = artifact_flow()
    resolved = resolve_references(flow, {"dataset": f"@run:{run_id}.dataset"})

    assert "study.dataset" in seed_values(flow, resolved)


def test_an_output_a_run_never_made_says_what_it_did(made_artifact):
    run_id, _reference = made_artifact
    with pytest.raises(RunRejected, match="prepare.dataset"):
        resolve_references(artifact_flow(), {"dataset": f"@run:{run_id}.weights"})


def test_a_reference_to_no_run_at_all_is_refused():
    with pytest.raises(RunRejected, match="no run"):
        resolve_references(artifact_flow(), {"dataset": "@run:nothing.dataset"})


def test_an_unknown_digest_is_refused():
    with pytest.raises(RunRejected, match="nothing here"):
        resolve_references(artifact_flow(), {"dataset": "sha256:" + "b2" * 32})


def test_a_reference_passed_whole_is_left_alone(made_artifact):
    """A python caller already has the object, and hands it over as one."""
    _run_id, reference = made_artifact
    assert resolve_references(artifact_flow(), {"dataset": reference}) == {
        "dataset": reference
    }