school_compare/backend/data_loader.py

"""
Data loading module — reads from marts.* tables built by dbt.
Provides efficient queries with caching.
"""

import logging
import re

import pandas as pd
import numpy as np
from typing import Optional, Dict, Tuple, List
import requests
from sqlalchemy import text
import sqlalchemy.exc
from sqlalchemy.orm import Session

from .config import settings
from .database import SessionLocal, engine
from .models import (
    DimSchool, DimLocation, KS2Performance,
    FactOfstedInspection, FactAdmissions,
    FactDeprivation, FactFinance, FactPupilCharacteristics,
)
from .ofsted_codes import ofsted_page_url, report_card_labels
from .schemas import SCHOOL_TYPE_MAP
from .gias_codes import (
    ADMISSIONS_POLICY,
    ESTABLISHMENT_STATUS,
    PHASE_OF_EDUCATION,
    RELIGIOUS_CHARACTER,
    SCHOOL_TYPE,
    translate,
)

# mart code column -> (API name column, dictionary)
_GIAS_CODE_COLUMNS = {
    "phase_code": ("phase", PHASE_OF_EDUCATION),
    "school_type_code": ("school_type", SCHOOL_TYPE),
    "status_code": ("status", ESTABLISHMENT_STATUS),
    "religious_character_code": ("religious_denomination", RELIGIOUS_CHARACTER),
    "admissions_policy_code": ("admissions_policy", ADMISSIONS_POLICY),
}


def translate_gias_code_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Map GIAS code columns to today's name columns (API contract).

    Runs immediately after pd.read_sql so every downstream consumer —
    filters, PHASE_GROUPS, payloads, /api/filters — keeps seeing names.
    DataFrames without the code columns (old schema, test fixtures) pass
    through unchanged.
    """
    for code_col, (name_col, mapping) in _GIAS_CODE_COLUMNS.items():
        if code_col in df.columns:
            df[name_col] = df[code_col].map(lambda c: translate(c, mapping))
    return df


_postcode_cache: Dict[str, Tuple[float, float]] = {}
_typesense_client = None


def _get_typesense_client():
    global _typesense_client
    if _typesense_client is not None:
        return _typesense_client
    url = settings.typesense_url
    key = settings.typesense_api_key
    if not url or not key:
        return None
    try:
        import typesense
        host = url.split("//")[-1]
        host_part, _, port_str = host.partition(":")
        port = int(port_str) if port_str else 8108
        _typesense_client = typesense.Client({
            "nodes": [{"host": host_part, "port": str(port), "protocol": "http"}],
            "api_key": key,
            "connection_timeout_seconds": 2,
        })
        return _typesense_client
    except Exception:
        return None


def search_schools_typesense(query: str, limit: int = 250) -> List[int]:
    """Search Typesense. Returns URNs in relevance order, or [] if unavailable."""
    client = _get_typesense_client()
    if client is None:
        return []
    try:
        result = client.collections["schools"].documents.search({
            "q": query,
            "query_by": "school_name,local_authority,postcode",
            "per_page": min(limit, 250),
            "typo_tokens_threshold": 1,
        })
        return [int(h["document"]["urn"]) for h in result.get("hits", [])]
    except Exception:
        return []


def normalize_school_type(school_type: Optional[str]) -> Optional[str]:
    """Convert cryptic school type codes to user-friendly names."""
    if not school_type:
        return None
    code = school_type.strip().upper()
    if code in SCHOOL_TYPE_MAP:
        return SCHOOL_TYPE_MAP[code]
    return school_type


def geocode_single_postcode(postcode: str) -> Optional[Tuple[float, float]]:
    """Geocode a single postcode using postcodes.io API."""
    if not postcode:
        return None
    postcode = postcode.strip().upper()
    if postcode in _postcode_cache:
        return _postcode_cache[postcode]
    try:
        response = requests.get(
            f"https://api.postcodes.io/postcodes/{postcode}",
            timeout=10,
        )
        if response.status_code == 200:
            data = response.json()
            if data.get("result"):
                lat = data["result"].get("latitude")
                lon = data["result"].get("longitude")
                if lat and lon:
                    _postcode_cache[postcode] = (lat, lon)
                    return (lat, lon)
    except Exception:
        pass
    return None


def haversine_distance(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Calculate great-circle distance between two points (miles)."""
    from math import radians, cos, sin, asin, sqrt
    lat1, lon1, lat2, lon2 = map(radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
    return 2 * asin(sqrt(a)) * 3956


# =============================================================================
# MAIN DATA LOAD — joins dim_school + dim_location + fact_performance
# fact_performance is a merged KS2+KS4 table (one row per URN per year).
# All-through schools have both KS2 and KS4 columns populated in the same row.
# =============================================================================

_MAIN_QUERY = text("""
    SELECT
        s.urn,
        s.school_name,
        s.phase_code,
        s.school_type_code,
        s.academy_trust_name   AS trust_name,
        s.academy_trust_uid    AS trust_uid,
        s.religious_character_code,
        s.gender,
        s.age_range,
        s.has_sixth_form,
        s.status_code,
        s.admissions_policy_code,
        s.capacity,
        s.total_pupils         AS gias_total_pupils,
        s.headteacher_name,
        s.website,
        foi.ofsted_grade,
        foi.ofsted_date,
        foi.ofsted_framework,
        l.local_authority_name AS local_authority,
        l.local_authority_code,
        l.address_line1        AS address1,
        l.address_line2        AS address2,
        l.town,
        l.postcode,
        l.latitude,
        l.longitude,
        p.year,
        p.source_urn,
        p.total_pupils,
        p.eligible_pupils,
        -- KS2 columns (NULL for pure secondary schools)
        p.rwm_expected_pct,
        p.rwm_high_pct,
        p.reading_expected_pct,
        p.reading_high_pct,
        p.reading_avg_score,
        p.reading_progress,
        p.reading_progress_lower_ci,
        p.reading_progress_upper_ci,
        p.writing_expected_pct,
        p.writing_high_pct,
        p.writing_progress,
        p.writing_progress_lower_ci,
        p.writing_progress_upper_ci,
        p.writing_working_towards_pct,
        p.maths_expected_pct,
        p.maths_high_pct,
        p.maths_avg_score,
        p.maths_progress,
        p.maths_progress_lower_ci,
        p.maths_progress_upper_ci,
        p.gps_expected_pct,
        p.gps_high_pct,
        p.gps_avg_score,
        p.science_expected_pct,
        p.reading_absence_pct,
        p.writing_absence_pct,
        p.maths_absence_pct,
        p.gps_absence_pct,
        p.science_absence_pct,
        p.rwm_expected_boys_pct,
        p.rwm_high_boys_pct,
        p.rwm_expected_girls_pct,
        p.rwm_high_girls_pct,
        p.rwm_expected_disadvantaged_pct,
        p.rwm_expected_non_disadvantaged_pct,
        p.disadvantaged_gap,
        p.disadvantaged_pct,
        p.eal_pct,
        p.stability_pct,
        -- KS4 columns (NULL for pure primary schools)
        p.attainment_8_score,
        p.progress_8_score,
        p.progress_8_lower_ci,
        p.progress_8_upper_ci,
        p.progress_8_english,
        p.progress_8_maths,
        p.progress_8_ebacc,
        p.progress_8_open,
        p.progress_8_banding,
        p.attainment_8_disadvantage_gap,
        p.progress_8_disadvantage_gap,
        p.english_maths_strong_pass_pct,
        p.english_maths_standard_pass_pct,
        p.ebacc_entry_pct,
        p.ebacc_strong_pass_pct,
        p.ebacc_standard_pass_pct,
        p.ebacc_avg_score,
        p.gcse_grade_91_pct,
        p.prior_attainment_avg,
        -- SEN (coalesced KS2+KS4 in fact_performance)
        p.sen_support_pct,
        p.sen_ehcp_pct
    FROM marts.dim_school s
    JOIN marts.dim_location l ON s.urn = l.urn
    LEFT JOIN marts.fact_performance p ON s.urn = p.urn
    LEFT JOIN (
        SELECT DISTINCT ON (urn)
            urn,
            -- Fall back to the ungraded-inspection grade when no graded grade exists.
            COALESCE(overall_effectiveness, ungraded_grade) AS ofsted_grade,
            inspection_date       AS ofsted_date,
            framework             AS ofsted_framework
        FROM marts.fact_ofsted_inspection
        ORDER BY urn, inspection_date DESC NULLS LAST
    ) foi ON s.urn = foi.urn
    ORDER BY s.school_name, p.year
""")

# Fallback used when marts.dim_school predates the has_sixth_form column
# (i.e. the nightly dbt pipeline hasn't rebuilt the mart yet on this DB).
# Keeps the column present as NULL so downstream code — including the
# app.py fallback branch — behaves as designed instead of KeyError-ing.
_MAIN_QUERY_NO_SIXTH_FORM = text(
    str(_MAIN_QUERY).replace("s.has_sixth_form,", "NULL AS has_sixth_form,")
)
assert "NULL AS has_sixth_form" in str(_MAIN_QUERY_NO_SIXTH_FORM), (
    "expected replacement of 's.has_sixth_form,' to have taken effect"
)

# Fallback used when marts.dim_school predates the GIAS code-dictionary
# migration (i.e. the nightly dbt pipeline hasn't rebuilt the mart yet on
# this DB, so it still has the old name columns instead of *_code columns).
_MAIN_QUERY_LEGACY_NAMES = str(_MAIN_QUERY)
_LEGACY_NAME_REPLACEMENTS = [
    ("s.phase_code,", "s.phase,"),
    ("s.school_type_code,", "s.school_type,"),
    (
        "s.religious_character_code,",
        "s.religious_character AS religious_denomination,",
    ),
    ("s.status_code,", "s.status,"),
    ("s.admissions_policy_code,", "s.admissions_policy,"),
]
for _old, _new in _LEGACY_NAME_REPLACEMENTS:
    assert _old in _MAIN_QUERY_LEGACY_NAMES, (
        f"expected {_old!r} to be present in _MAIN_QUERY before replacement"
    )
    _MAIN_QUERY_LEGACY_NAMES = _MAIN_QUERY_LEGACY_NAMES.replace(_old, _new)
_MAIN_QUERY_LEGACY_NAMES = text(_MAIN_QUERY_LEGACY_NAMES)

_GIAS_CODE_COLUMN_NAMES = (
    "phase_code",
    "school_type_code",
    "religious_character_code",
    "status_code",
    "admissions_policy_code",
)

_MISSING_COLUMN_RE = re.compile(r'column "?(?:s\.)?(\w+)"? does not exist')


def _missing_column_name(exc: Exception) -> Optional[str]:
    """Name of the missing column from a psycopg2 UndefinedColumn error.

    Inspects exc.orig (the DBAPI error), whose message names only the
    offending column — str(exc) also embeds the full SQL statement, which
    contains every column name and therefore must not be matched against.
    """
    orig = getattr(exc, "orig", None)
    match = _MISSING_COLUMN_RE.search(str(orig) if orig is not None else str(exc))
    return match.group(1) if match else None


def load_school_data_as_dataframe() -> pd.DataFrame:
    """Load all school + KS2 data as a pandas DataFrame."""
    try:
        df = pd.read_sql(_MAIN_QUERY, engine)
    except sqlalchemy.exc.ProgrammingError as exc:
        missing = _missing_column_name(exc)
        if missing in _GIAS_CODE_COLUMN_NAMES:
            logging.getLogger(__name__).warning(
                "marts predate the GIAS code migration — falling back to "
                "legacy name-column query: %s",
                exc,
            )
            try:
                df = pd.read_sql(_MAIN_QUERY_LEGACY_NAMES, engine)
            except Exception as exc2:
                print(f"Warning: Could not load school data from marts: {exc2}")
                return pd.DataFrame()
        elif missing == "has_sixth_form":
            logging.getLogger(__name__).warning(
                "marts.dim_school is missing has_sixth_form (pipeline hasn't "
                "rebuilt the mart yet on this DB) — retrying without it: %s",
                exc,
            )
            try:
                df = pd.read_sql(_MAIN_QUERY_NO_SIXTH_FORM, engine)
            except Exception as exc2:
                print(f"Warning: Could not load school data from marts: {exc2}")
                return pd.DataFrame()
        else:
            print(f"Warning: Could not load school data from marts: {exc}")
            return pd.DataFrame()
    except Exception as exc:
        print(f"Warning: Could not load school data from marts: {exc}")
        return pd.DataFrame()

    if df.empty:
        return df

    df = translate_gias_code_columns(df)

    # Build address string
    df["address"] = df.apply(
        lambda r: ", ".join(
            p for p in [r.get("address1"), r.get("address2"), r.get("town"), r.get("postcode")]
            if p and str(p) != "None"
        ),
        axis=1,
    )

    # Normalize school type
    df["school_type"] = df["school_type"].apply(normalize_school_type)

    return df


# Cache for DataFrame
_df_cache: Optional[pd.DataFrame] = None
# Pre-computed latest-year snapshot (one row per school, with prev-year trend columns)
_df_latest_cache: Optional[pd.DataFrame] = None


def load_school_data() -> pd.DataFrame:
    """Load school data with caching."""
    global _df_cache
    if _df_cache is not None:
        return _df_cache
    print("Loading school data from marts...")
    _df_cache = load_school_data_as_dataframe()
    if not _df_cache.empty:
        print(f"Total records loaded: {len(_df_cache)}")
        print(f"Unique schools: {_df_cache['urn'].nunique()}")
        print(f"Years: {sorted(_df_cache['year'].dropna().unique())}")
    else:
        print("No data found in marts (EES data may not have been loaded yet)")
    return _df_cache


def load_latest_school_data() -> pd.DataFrame:
    """Return a cached one-row-per-school DataFrame at the latest available year.

    The expensive groupby / merge / prev-year trend computation runs once at
    startup (or after a cache clear) rather than on every search request.
    Per-request filters (phase, gender, LA …) should be applied to the returned
    DataFrame's copy; they must NOT modify the cached object.
    """
    global _df_latest_cache
    if _df_latest_cache is not None:
        return _df_latest_cache

    df = load_school_data()
    if df.empty:
        return df

    # Schools that have no performance rows (PRUs, new schools, etc.)
    df_no_perf = df[df["year"].isna()].drop_duplicates(subset=["urn"])
    df_with_perf = df[df["year"].notna()]

    # Reduce to the latest year per school
    latest_year = df_with_perf.groupby("urn")["year"].max().reset_index()
    df_latest = df_with_perf.merge(latest_year, on=["urn", "year"])

    # Attach previous-year metrics for trend arrows (second-latest year per school)
    df_sorted = df_with_perf.sort_values(["urn", "year"], ascending=[True, False])
    df_prev = df_sorted.groupby("urn").nth(1).reset_index()
    if not df_prev.empty and "rwm_expected_pct" in df_prev.columns:
        prev_rwm = df_prev[["urn", "rwm_expected_pct"]].rename(
            columns={"rwm_expected_pct": "prev_rwm_expected_pct"}
        )
        if "attainment_8_score" in df_prev.columns:
            prev_rwm = prev_rwm.merge(
                df_prev[["urn", "attainment_8_score"]].rename(
                    columns={"attainment_8_score": "prev_attainment_8_score"}
                ),
                on="urn",
                how="outer",
            )
        df_latest = df_latest.merge(prev_rwm, on="urn", how="left")

    # Merge back schools with no performance data
    df_latest = pd.concat([df_latest, df_no_perf], ignore_index=True)

    print(f"Latest-snapshot cache built: {len(df_latest)} schools")
    _df_latest_cache = df_latest
    return _df_latest_cache


def clear_cache():
    """Clear all caches."""
    global _df_cache, _df_latest_cache
    _df_cache = None
    _df_latest_cache = None


# =============================================================================
# METADATA QUERIES
# =============================================================================

def get_available_years(db: Session = None) -> List[int]:
    close_db = db is None
    if db is None:
        db = SessionLocal()
    try:
        result = db.query(KS2Performance.year).distinct().order_by(KS2Performance.year).all()
        return [r[0] for r in result]
    except Exception:
        return []
    finally:
        if close_db:
            db.close()


def get_available_local_authorities(db: Session = None) -> List[str]:
    close_db = db is None
    if db is None:
        db = SessionLocal()
    try:
        result = (
            db.query(DimLocation.local_authority_name)
            .filter(DimLocation.local_authority_name.isnot(None))
            .distinct()
            .order_by(DimLocation.local_authority_name)
            .all()
        )
        return [r[0] for r in result if r[0]]
    except Exception:
        return []
    finally:
        if close_db:
            db.close()


def get_schools_count(db: Session = None) -> int:
    close_db = db is None
    if db is None:
        db = SessionLocal()
    try:
        return db.query(DimSchool).count()
    except Exception:
        return 0
    finally:
        if close_db:
            db.close()


def get_data_info(db: Session = None) -> dict:
    close_db = db is None
    if db is None:
        db = SessionLocal()
    try:
        school_count = get_schools_count(db)
        years = get_available_years(db)
        local_authorities = get_available_local_authorities(db)
        return {
            "total_schools": school_count,
            "years_available": years,
            "local_authorities_count": len(local_authorities),
            "data_source": "PostgreSQL (marts)",
        }
    finally:
        if close_db:
            db.close()


# =============================================================================
# SUPPLEMENTARY DATA — per-school detail page
# =============================================================================

def compute_benchmarks(df: pd.DataFrame, census_benchmarks: dict | None = None) -> dict:
    """State-school benchmarks computed from our dataset (spec §5/§8.6).

    NOT official DfE figures — consumers must label them
    "state-school average (computed from our dataset)". The disadvantaged
    attainment average is weighted by cohort size (eligible_pupils) so
    small schools don't dominate.

    Context measures (FSM/EAL/pupil counts) come from `census_benchmarks`
    (the fact_census_benchmarks mart, pupil-weighted, keyed by phase): the
    performance df has no fsm_pct at all, and its eal/disadvantaged columns
    are KS2-only — medianing them for "secondary" produced junk anchors
    from the handful of all-through schools. When the mart is unavailable
    these are None; never fall back across measure definitions.
    """
    if df.empty or "year" not in df.columns:
        return {}
    latest_year = df["year"].max()
    if pd.isna(latest_year):
        return {}
    d = df[df["year"] == latest_year]
    if d.empty:
        return {}
    is_secondary = (
        d["attainment_8_score"].notna()
        if "attainment_8_score" in d.columns
        else pd.Series(False, index=d.index)
    )
    prim, sec = d[~is_secondary], d[is_secondary]

    def _median(sub, col):
        if col not in sub.columns:
            return None
        v = sub[col].median()
        return round(float(v), 1) if pd.notna(v) else None

    def _weighted_disadvantaged(sub):
        needed = {"rwm_expected_disadvantaged_pct", "eligible_pupils"}
        if not needed <= set(sub.columns):
            return None
        s = sub.dropna(subset=list(needed))
        if s.empty or s["eligible_pupils"].sum() == 0:
            return None
        w = (
            (s["rwm_expected_disadvantaged_pct"] * s["eligible_pupils"]).sum()
            / s["eligible_pupils"].sum()
        )
        return round(float(w), 1)

    def _block(sub, phase, with_disadvantaged):
        census = (census_benchmarks or {}).get(phase) or {}
        block = {
            "eal_pct": census.get("eal_pct"),
            "sen_support_pct": _median(sub, "sen_support_pct"),
            "disadvantaged_pct": _median(sub, "disadvantaged_pct") if with_disadvantaged else None,
            "fsm_pct": census.get("fsm_pct"),
            "median_pupils": census.get("median_pupils"),
        }
        if with_disadvantaged:
            block["disadvantaged_rwm_expected_pct"] = _weighted_disadvantaged(sub)
        return block

    return {
        "source": "state-school average (computed from our dataset)",
        "year": int(latest_year),
        "primary": _block(prim, "primary", with_disadvantaged=True),
        "secondary": _block(sec, "secondary", with_disadvantaged=False),
    }


def _ofsted_block(o, urn: int) -> dict:
    """Serialize the latest Ofsted inspection row for API responses.

    `grade_source` records where the effective overall grade came from:
    a graded (Section 5) inspection, or carried forward from an ungraded
    (Section 8) outcome — materially different claims a UI must be able
    to distinguish. `report_card` holds coded+labelled renewed-framework
    (Nov 2025) area judgements; safeguarding is a separate boolean and
    never appears among the graded areas.
    """
    if o.overall_effectiveness is not None:
        grade_source = "graded"
        overall = o.overall_effectiveness
    elif o.ungraded_grade is not None:
        # Fall back to the grade parsed from an ungraded (Section 8) outcome
        # (e.g. "School remains Good") so the detail page matches the list badge.
        grade_source = "ungraded_carried_forward"
        overall = o.ungraded_grade
    else:
        grade_source = None
        overall = None

    block = {
        "framework": o.framework,
        "inspection_date": o.inspection_date.isoformat() if o.inspection_date else None,
        "rc_inspection_date": (
            o.rc_inspection_date.isoformat()
            if getattr(o, "rc_inspection_date", None)
            else None
        ),
        "inspection_type": o.inspection_type,
        "overall_effectiveness": overall,
        "grade_source": grade_source,
        "quality_of_education": o.quality_of_education,
        "behaviour_attitudes": o.behaviour_attitudes,
        "personal_development": o.personal_development,
        "leadership_management": o.leadership_management,
        "early_years_provision": o.early_years_provision,
        "sixth_form_provision": o.sixth_form_provision,
        "previous_overall": None,  # Not available in new schema
        "rc_safeguarding_met": o.rc_safeguarding_met,
        "rc_inclusion": o.rc_inclusion,
        "rc_curriculum_teaching": o.rc_curriculum_teaching,
        "rc_achievement": o.rc_achievement,
        "rc_attendance_behaviour": o.rc_attendance_behaviour,
        "rc_personal_development": o.rc_personal_development,
        "rc_leadership_governance": o.rc_leadership_governance,
        "rc_early_years": o.rc_early_years,
        "rc_sixth_form": o.rc_sixth_form,
        "report_url": o.report_url,
        "ofsted_page_url": ofsted_page_url(urn),
    }
    block["report_card"] = report_card_labels(block)
    return block


def _admissions_row_dict(a) -> dict:
    """Serialize one fact_admissions row for API responses."""
    return {
        "year": a.year,
        "school_phase": a.school_phase,
        "places_offered": a.places_offered,
        "total_applications": a.total_applications,
        "first_preference_applications": a.first_preference_applications,
        "first_preference_offers": a.first_preference_offers,
        "first_preference_offer_pct": a.first_preference_offer_pct,
        "oversubscription_ratio": a.oversubscription_ratio,
        "oversubscribed": a.oversubscribed,
        "total_offers": a.total_offers,
        "second_preference_offers": a.second_preference_offers,
        "third_preference_offers": a.third_preference_offers,
        "cross_la_applications": a.cross_la_applications,
        "cross_la_offers": a.cross_la_offers,
    }


def _census_dict(pc) -> dict:
    return {
        "year": pc.year,
        "total_pupils": pc.total_pupils,
        "female_pupils": pc.female_pupils,
        "male_pupils": pc.male_pupils,
        "fsm_pct": pc.fsm_pct,
        "eal_pct": pc.eal_pct,
    }


def _deprivation_dict(d) -> dict:
    return {
        "lsoa_code": d.lsoa_code,
        "idaci_score": d.idaci_score,
        "idaci_decile": d.idaci_decile,
    }


def _finance_dict(f) -> dict:
    return {
        "year": f.year,
        "per_pupil_spend": f.per_pupil_spend,
        "staff_cost_pct": f.staff_cost_pct,
        "teacher_cost_pct": f.teacher_cost_pct,
        "support_staff_cost_pct": f.support_staff_cost_pct,
        "premises_cost_pct": f.premises_cost_pct,
    }


def _empty_supplementary() -> dict:
    return {
        "ofsted": None,
        "census": None,
        "admissions": None,
        "admissions_history": [],
        "sen_detail": None,
        "phonics": None,
        "deprivation": None,
        "finance": None,
    }


def get_supplementary_data_batch(db: Session, urns: list[int]) -> dict:
    """Fetch supplementary data for many URNs with one query per table
    (WHERE urn IN (...)) instead of ~5 queries per school, collapsing the
    per-request round-trips from 5*N to a constant 5. Returns {urn: block}
    with the same shape get_supplementary_data produces per URN.

    Each table is queried independently and failures degrade that table to
    empty for every URN — a missing mart never blanks the others.
    """
    urns = [int(u) for u in urns]
    result = {urn: _empty_supplementary() for urn in urns}
    if not urns:
        return result

    def _safe(fn):
        try:
            fn()
        except Exception as e:
            import logging
            logging.getLogger(__name__).error("batch supplementary query failed: %s", e)
            db.rollback()

    # Ofsted — latest inspection per URN. Ordered so the first row seen per
    # URN is the most recent.
    def _ofsted():
        rows = (
            db.query(FactOfstedInspection)
            .filter(FactOfstedInspection.urn.in_(urns))
            .order_by(FactOfstedInspection.urn, FactOfstedInspection.inspection_date.desc())
            .all()
        )
        seen = set()
        for o in rows:
            if o.urn in seen:
                continue
            seen.add(o.urn)
            result[o.urn]["ofsted"] = _ofsted_block(o, o.urn)
    _safe(_ofsted)

    # Census — latest year per URN.
    def _census():
        rows = (
            db.query(FactPupilCharacteristics)
            .filter(FactPupilCharacteristics.urn.in_(urns))
            .order_by(FactPupilCharacteristics.urn, FactPupilCharacteristics.year.desc())
            .all()
        )
        seen = set()
        for pc in rows:
            if pc.urn in seen:
                continue
            seen.add(pc.urn)
            result[pc.urn]["census"] = _census_dict(pc)
    _safe(_census)

    # Admissions — all years per URN, oldest first (multi-year trend view).
    def _admissions():
        rows = (
            db.query(FactAdmissions)
            .filter(FactAdmissions.urn.in_(urns))
            .order_by(FactAdmissions.urn, FactAdmissions.year.asc())
            .all()
        )
        history: dict = {urn: [] for urn in urns}
        for a in rows:
            history[a.urn].append(_admissions_row_dict(a))
        for urn, rows_for_urn in history.items():
            result[urn]["admissions_history"] = rows_for_urn
            result[urn]["admissions"] = rows_for_urn[-1] if rows_for_urn else None
    _safe(_admissions)

    # Deprivation — one row per URN.
    def _deprivation():
        rows = (
            db.query(FactDeprivation)
            .filter(FactDeprivation.urn.in_(urns))
            .all()
        )
        for d in rows:
            result[d.urn]["deprivation"] = _deprivation_dict(d)
    _safe(_deprivation)

    # Finance — latest year per URN.
    def _finance():
        rows = (
            db.query(FactFinance)
            .filter(FactFinance.urn.in_(urns))
            .order_by(FactFinance.urn, FactFinance.year.desc())
            .all()
        )
        seen = set()
        for f in rows:
            if f.urn in seen:
                continue
            seen.add(f.urn)
            result[f.urn]["finance"] = _finance_dict(f)
    _safe(_finance)

    return result


def get_supplementary_data(db: Session, urn: int) -> dict:
    """Supplementary data for a single URN (thin wrapper over the batch)."""
    return get_supplementary_data_batch(db, [urn])[int(urn)]