"""API-boundary translation: marts now carry GIAS codes; the DataFrame the
rest of the backend sees must carry today's name strings."""

import numpy as np
import pandas as pd

from backend.data_loader import translate_gias_code_columns
from backend.gias_codes import ESTABLISHMENT_STATUS, PHASE_OF_EDUCATION


def _code_for(mapping, name):
    return next(c for c, n in mapping.items() if n == name)


def test_codes_become_todays_names():
    df = pd.DataFrame([{
        "urn": 1,
        "phase_code": float(_code_for(PHASE_OF_EDUCATION, "Primary")),
        "school_type_code": np.nan,
        "status_code": float(_code_for(ESTABLISHMENT_STATUS, "Open, but proposed to close")),
        "religious_character_code": np.nan,
        "admissions_policy_code": np.nan,
    }])
    out = translate_gias_code_columns(df)
    row = out.iloc[0]
    assert row["phase"] == "Primary"
    assert row["status"] == "Open, but proposed to close"
    assert row["school_type"] is None
    assert row["religious_denomination"] is None
    assert row["admissions_policy"] is None


def test_unknown_code_degrades_not_blanks():
    df = pd.DataFrame([{"urn": 1, "phase_code": 9999.0}])
    out = translate_gias_code_columns(df)
    assert out.iloc[0]["phase"] == "Unknown (9999)"


def test_missing_code_columns_are_a_noop():
    """Old-schema DataFrames (tests, pre-pipeline DBs) pass through untouched."""
    df = pd.DataFrame([{"urn": 1, "phase": "Primary", "status": "Open"}])
    out = translate_gias_code_columns(df)
    assert out.iloc[0]["phase"] == "Primary"
    assert out.iloc[0]["status"] == "Open"


def test_load_school_data_survives_premigration_marts(monkeypatch):
    """Real prod state until the nightly pipeline first rebuilds the mart with
    the GIAS code columns: marts.dim_school still has the old name columns
    (phase, school_type, religious_character, status, admissions_policy)
    instead of the new *_code columns. The first query raises UndefinedColumn
    on s.phase_code; load_school_data_as_dataframe must retry with the
    legacy name-column query rather than swallow the error and return (and
    then have load_school_data cache) an empty DataFrame."""
    import sqlalchemy.exc
    from backend import data_loader

    data_loader._df_cache = None
    data_loader._df_latest_cache = None

    good_df = pd.DataFrame(
        [
            {
                "urn": 1,
                "school_name": "Legacy School",
                "phase": "Primary",
                "school_type": "Academy",
                "status": "Open",
            }
        ]
    )
    calls = []

    def fake_read_sql(query, con):
        calls.append(query)
        if len(calls) == 1:
            raise sqlalchemy.exc.ProgrammingError(
                "(psycopg2.errors.UndefinedColumn) column s.phase_code does not exist",
                None,
                None,
            )
        return good_df.copy()

    monkeypatch.setattr(data_loader.pd, "read_sql", fake_read_sql)

    try:
        df = data_loader.load_school_data_as_dataframe()
    finally:
        data_loader._df_cache = None
        data_loader._df_latest_cache = None

    assert len(calls) == 2, "must retry with the legacy name-column query variant"
    assert calls[1] is data_loader._MAIN_QUERY_LEGACY_NAMES
    assert not df.empty
    assert df["phase"].iloc[0] == "Primary"
    assert df["status"].iloc[0] == "Open"