"""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, col, select

from fluksio.core.config import settings
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, RunMetric, 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)
    hit = cache.lookup(plain)
    assert hit is not None and hit.outputs == {"study.loss": 1.5}
    # Where its series is, so a run reusing it can read the curve back.
    assert hit.metrics_run == "cache-1"
    found = cache.lookup(with_artifact)
    assert found is not None and found.outputs == {"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) is None
    assert cache.lookup("never-seen") is None
    assert cache.lookup("") is 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
    }


def chaining_flow() -> FlowDef:
    """A flow taking a json config and a label another run worked out."""
    meta = MessageSpec(name="meta", dtype=DType.JSON)
    label = MessageSpec(name="label", dtype=DType.STR)
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=meta), FlowInput(spec=label)],
        nodes=[NodeDef(id="train", requires=[meta, label])],
    )


@pytest.fixture
def made_config():
    """A finished run whose result is an object, not bytes."""
    meta = {"rows": 256, "source": "builtin"}
    run_id = new_run_id()
    with Session(db_engine) as session:
        session.add(
            Run(
                id=run_id,
                flow="generate",
                status="ok",
                result={"meta": meta, "label": "run-7"},
                created_at=datetime.now(UTC),
            )
        )
        session.commit()
    yield run_id, meta
    with Session(db_engine) as session:
        session.delete(session.get(Run, run_id))
        session.commit()


def test_a_json_input_may_name_a_run_s_output(made_config):
    """The gap this closes: chaining without pasting the object into a shell."""
    run_id, meta = made_config
    flow = chaining_flow()
    resolved = resolve_references(flow, {"meta": f"@run:{run_id}.meta"})

    assert resolved["meta"] == meta
    # And it is the value's own type from here on, so the input check passes.
    assert seed_values(flow, resolved)["study.meta"] == meta


def test_the_spelling_is_reserved_on_a_text_input_too(made_config):
    run_id, _meta = made_config
    resolved = resolve_references(chaining_flow(), {"label": f"@run:{run_id}.label"})

    assert resolved["label"] == "run-7"


def test_text_that_names_nothing_is_still_left_alone(made_config):
    """Only the two spellings are read as names; everything else is a value."""
    _run_id, _meta = made_config
    params = {"label": "@run-of-the-mill", "meta": {"rows": 1}}
    assert resolve_references(chaining_flow(), params) == params


def test_the_overview_counts_a_flow_the_list_page_would_not_reach(
    client, superuser_token_headers
):
    """The list caps at 500 newest; the flow rail needs whole counts."""
    with Session(db_engine) as session:
        for index in range(3):
            session.add(
                Run(
                    id=f"ov-{index}",
                    flow="overviewed",
                    status="ok" if index else "running",
                    created_at=datetime(2026, 1, 1 + index, tzinfo=UTC),
                )
            )
        session.commit()

    rows = client.get(
        f"{settings.API_V1_STR}/runs/overview", headers=superuser_token_headers
    ).json()
    row = next(r for r in rows if r["flow"] == "overviewed")

    assert (row["runs"], row["running"], row["queued"]) == (3, 1, 0)


def test_overview_is_not_read_as_a_run_id(client, superuser_token_headers):
    """`/overview` is declared before `/{run_id}`, which would swallow it."""
    answer = client.get(
        f"{settings.API_V1_STR}/runs/overview", headers=superuser_token_headers
    )

    assert answer.status_code == 200
    assert isinstance(answer.json(), list)


# -----------------------------------------------------------------------------
# A cached node's curve
#
# A hit replays no emissions, so the series stays in the run that recorded it
# and the run reusing it points there. Reading either one answers the same.
# -----------------------------------------------------------------------------


@pytest.fixture
def reused_run():
    """A run of `quick` whose node was restored from a run of `train`."""
    with Session(db_engine) as session:
        made = datetime.now(UTC)
        session.add(Run(id="src-1", flow="train", status="ok", created_at=made))
        session.add(Run(id="reuse-1", flow="quick", status="ok", created_at=made))
        session.add(RunNode(run_id="src-1", node="train.fit", status="ok"))
        session.add(
            RunNode(
                run_id="reuse-1",
                node="quick.fit",
                status="cached",
                cached_from="src-1",
            )
        )
        for step, value in enumerate([3.0, 2.0, 1.0]):
            session.add(
                RunMetric(
                    run_id="src-1",
                    name="train.loss",
                    step=step,
                    node="train.fit",
                    value=value,
                )
            )
        session.commit()
    yield
    with Session(db_engine) as session:
        for row in session.exec(select(RunMetric)).all():
            session.delete(row)
        for row in session.exec(select(RunNode)).all():
            session.delete(row)
        for run_id in ("src-1", "reuse-1"):
            run = session.get(Run, run_id)
            if run is not None:
                session.delete(run)
        session.commit()


@pytest.mark.usefixtures("reused_run")
def test_a_cached_node_answers_with_the_curve_it_was_restored_from(
    client, superuser_token_headers
):
    points = client.get(
        f"{settings.API_V1_STR}/runs/reuse-1/metrics",
        headers=superuser_token_headers,
    ).json()

    # Named for the flow that asked, not the one that recorded it.
    assert [point["name"] for point in points] == ["quick.loss"] * 3
    assert [point["value"] for point in points] == [3.0, 2.0, 1.0]

    named = client.get(
        f"{settings.API_V1_STR}/runs/reuse-1/metrics",
        params={"name": "quick.loss"},
        headers=superuser_token_headers,
    ).json()
    assert len(named) == 3


@pytest.mark.usefixtures("reused_run")
def test_a_curve_whose_run_is_gone_is_empty_rather_than_an_error(
    client, superuser_token_headers
):
    """Deleting a flow deletes its runs; what pointed at one is left holding it."""
    with Session(db_engine) as session:
        session.delete(session.get(Run, "src-1"))
        session.commit()

    answer = client.get(
        f"{settings.API_V1_STR}/runs/reuse-1/metrics",
        headers=superuser_token_headers,
    )
    assert answer.status_code == 200
    assert answer.json() == []


def _metric(run_id: str, name: str, step: int, value: float, ts: float) -> RunMetric:
    return RunMetric(run_id=run_id, name=name, step=step, value=value, ts=ts)


@pytest.fixture
def plotted():
    """A run with a loss curve and an epoch counter beside it."""
    run_id = new_run_id()
    with Session(db_engine) as session:
        session.add(
            Run(id=run_id, flow="study", status="ok", created_at=datetime.now(UTC))
        )
        for step, (loss, epoch) in enumerate([(1.0, 10.0), (0.5, 20.0), (0.25, 30.0)]):
            session.add(_metric(run_id, "study.loss", step, loss, 100.0 + step * 5))
            session.add(_metric(run_id, "study.epoch", step, epoch, 100.0 + step * 5))
        session.commit()
    yield run_id
    with Session(db_engine) as session:
        for row in session.exec(
            select(RunMetric).where(col(RunMetric.run_id) == run_id)
        ).all():
            session.delete(row)
        session.delete(session.get(Run, run_id))
        session.commit()


def test_a_comparison_is_plotted_against_the_step_by_default(
    client, superuser_token_headers, plotted
):
    answer = client.get(
        f"{settings.API_V1_STR}/runs/series/compare",
        params={"ids": plotted, "metric": "study.loss"},
        headers=superuser_token_headers,
    ).json()

    assert answer["x"] == "step"
    assert answer["lines"][0]["points"] == [[0.0, 1.0], [1.0, 0.5], [2.0, 0.25]]


def test_time_is_measured_from_this_runs_own_first_reading(
    client, superuser_token_headers, plotted
):
    """Runs started hours apart still lie on top of each other."""
    answer = client.get(
        f"{settings.API_V1_STR}/runs/series/compare",
        params={"ids": plotted, "metric": "study.loss", "x": "time"},
        headers=superuser_token_headers,
    ).json()

    assert answer["x"] == "time"
    assert answer["lines"][0]["points"] == [[0.0, 1.0], [5.0, 0.5], [10.0, 0.25]]


def test_one_metric_can_be_plotted_against_another(
    client, superuser_token_headers, plotted
):
    answer = client.get(
        f"{settings.API_V1_STR}/runs/series/compare",
        params={"ids": plotted, "metric": "study.loss", "x": "study.epoch"},
        headers=superuser_token_headers,
    ).json()

    assert answer["lines"][0]["points"] == [[10.0, 1.0], [20.0, 0.5], [30.0, 0.25]]


def test_a_step_the_x_metric_never_reached_is_left_out(
    client, superuser_token_headers, plotted
):
    """The join is on the step, which is the only thing two series share."""
    with Session(db_engine) as session:
        session.add(_metric(plotted, "study.loss", 3, 0.1, 120.0))
        session.commit()

    answer = client.get(
        f"{settings.API_V1_STR}/runs/series/compare",
        params={"ids": plotted, "metric": "study.loss", "x": "study.epoch"},
        headers=superuser_token_headers,
    ).json()

    assert [point[0] for point in answer["lines"][0]["points"]] == [10.0, 20.0, 30.0]