"""get_supplementary_data_batch fetches one query per table for all URNs
(not ~5 per school) and returns the same per-URN block shape as the
single-URN function, picking the latest row per URN where relevant."""

import types

from backend import data_loader
from backend.data_loader import get_supplementary_data_batch


class _FakeQuery:
    """Records that a query ran and serves canned rows filtered by an in-list."""

    def __init__(self, recorder, model_name, rows):
        self._rec = recorder
        self._model = model_name
        self._rows = rows

    def filter(self, *args, **kwargs):
        return self

    def order_by(self, *args, **kwargs):
        return self

    def all(self):
        self._rec.append(self._model)
        return self._rows

    def first(self):
        self._rec.append(self._model)
        return self._rows[0] if self._rows else None


class _FakeSession:
    def __init__(self, rows_by_model):
        self.rows_by_model = rows_by_model
        self.queries: list[str] = []

    def query(self, model):
        name = model.__name__
        return _FakeQuery(self.queries, name, self.rows_by_model.get(name, []))

    def rollback(self):
        pass


def _ofsted_row(urn, date, oe):
    base = {f: None for f in (
        "framework", "inspection_type", "quality_of_education", "behaviour_attitudes",
        "personal_development", "leadership_management", "early_years_provision",
        "sixth_form_provision", "ungraded_outcome", "ungraded_grade",
        "rc_safeguarding_met", "rc_inclusion", "rc_curriculum_teaching", "rc_achievement",
        "rc_attendance_behaviour", "rc_personal_development", "rc_leadership_governance",
        "rc_early_years", "rc_sixth_form", "report_url",
    )}
    base.update(urn=urn, inspection_date=types.SimpleNamespace(isoformat=lambda: date),
                overall_effectiveness=oe, grade_source=None)
    return types.SimpleNamespace(**base)


def _adm_row(urn, year):
    return types.SimpleNamespace(
        urn=urn, year=year, school_phase="Primary", places_offered=100,
        total_applications=200, first_preference_applications=150,
        first_preference_offers=140, first_preference_offer_pct=93.3,
        oversubscription_ratio=1.5, oversubscribed=True,
        total_offers=100, second_preference_offers=5, third_preference_offers=2,
        cross_la_applications=10, cross_la_offers=3,
    )


def test_one_query_per_table_and_latest_row_per_urn():
    rows = {
        # URN 1 has two Ofsted rows; the batch must keep the most recent (2023).
        "FactOfstedInspection": [
            _ofsted_row(1, "2023-01-01", 2),
            _ofsted_row(1, "2019-01-01", 3),
            _ofsted_row(2, "2021-06-01", 1),
        ],
        "FactAdmissions": [_adm_row(1, 202526), _adm_row(1, 202627), _adm_row(2, 202627)],
        "FactPupilCharacteristics": [],
        "FactDeprivation": [],
        "FactFinance": [],
    }
    session = _FakeSession(rows)
    out = get_supplementary_data_batch(session, [1, 2])

    # Exactly one query per table — five total, regardless of two URNs.
    assert sorted(session.queries) == [
        "FactAdmissions", "FactDeprivation", "FactFinance",
        "FactOfstedInspection", "FactPupilCharacteristics",
    ]

    # Latest Ofsted kept per URN
    assert out[1]["ofsted"]["overall_effectiveness"] == 2
    assert out[2]["ofsted"]["overall_effectiveness"] == 1

    # Admissions history grouped per URN, latest exposed as `admissions`
    assert [r["year"] for r in out[1]["admissions_history"]] == [202526, 202627]
    assert out[1]["admissions"]["year"] == 202627
    assert out[2]["admissions_history"] == [{**out[2]["admissions_history"][0]}]

    # Empty tables degrade to the null block, not a crash
    assert out[1]["census"] is None and out[1]["deprivation"] is None


def test_single_wrapper_matches_batch(monkeypatch):
    session = _FakeSession({"FactOfstedInspection": [_ofsted_row(5, "2022-01-01", 2)]})
    single = data_loader.get_supplementary_data(session, 5)
    assert single["ofsted"]["overall_effectiveness"] == 2
    assert single["admissions_history"] == []