"""_national_averages_payload reads persisted marts (computed at import
time) — it must never loop the dataframe per year. The only dataframe work
allowed is the single-latest-year KS4 fallback for the window between a
deploy and the next DAG run."""

import numpy as np
import pandas as pd
import pytest

LATEST = 202425


def _df():
    return pd.DataFrame(
        [
            dict(year=202324, attainment_8_score=40.0, rwm_expected_pct=np.nan),
            dict(year=LATEST, attainment_8_score=50.0, rwm_expected_pct=np.nan),
            dict(year=LATEST, attainment_8_score=30.0, rwm_expected_pct=np.nan),
            dict(year=LATEST, attainment_8_score=np.nan, rwm_expected_pct=80.0),
        ]
    )


class _Ks2Row:
    year = LATEST
    rwm_expected_pct = 62.1
    gps_expected_pct = 72.0


class _Ks4Row:
    year = LATEST
    attainment_8_score = 46.5
    progress_8_score = -0.02


class _StubSession:
    """Returns KS2 rows for the first query and KS4 rows for the second —
    mirroring the payload's query order."""

    def __init__(self):
        self.calls = 0

    def query(self, model):
        self._model = model.__name__
        return self

    def order_by(self, *a):
        return self

    def all(self):
        return [_Ks2Row()] if self._model == "Ks2NationalAverage" else [_Ks4Row()]

    def close(self):
        pass


class _Ks4MissingSession(_StubSession):
    def all(self):
        if self._model == "Ks4NationalAverage":
            raise RuntimeError("relation does not exist")
        return [_Ks2Row()]

    def rollback(self):
        pass


@pytest.fixture()
def payload(monkeypatch):
    from backend import app as app_module
    from backend import database as database_module

    def _run(session_cls):
        monkeypatch.setattr(database_module, "SessionLocal", session_cls)
        return app_module._national_averages_payload(_df())

    return _run


def test_ks4_averages_come_from_the_mart_not_the_dataframe(payload):
    body = payload(_StubSession)
    # Mart value (46.5), NOT the dataframe mean of (50+30)/2 = 40.0
    assert body["secondary"]["attainment_8_score"] == 46.5
    assert body["primary"]["rwm_expected_pct"] == 62.1
    assert body["by_year"][-1]["secondary"]["progress_8_score"] == -0.02


def test_missing_ks4_mart_falls_back_to_latest_year_only(payload):
    body = payload(_Ks4MissingSession)
    # Fallback computes the latest year from the df: mean(50, 30) = 40.0
    assert body["secondary"]["attainment_8_score"] == 40.0
    # ...and only the latest year — no historical KS4 loop
    ks4_years = [e["year"] for e in body["by_year"] if e["secondary"]]
    assert ks4_years == [LATEST]