import polars as pl
import pytest

from pipeline.download.education import (
    BAND_MAP,
    OUTPUT_BUCKETS,
    _qualification_percentages,
)


def _long_rows(geo: str, counts: dict[str, int]) -> list[dict]:
    """Build NOMIS-shaped long rows for one LSOA from {band_label: count}.

    NOMIS emits a 0-count row when an LSOA has none in a band, so bands not
    given default to 0 to mirror that.
    """
    return [
        {
            "GEOGRAPHY_CODE": geo,
            "C2021_HIQUAL_8_NAME": label,
            "OBS_VALUE": counts.get(label, 0),
        }
        for label in BAND_MAP
    ]


def test_qualification_percentages_keyed_by_lsoa_with_seven_buckets():
    df = pl.DataFrame(
        _long_rows(
            "E01000001",
            {
                "No qualifications": 10,
                "Level 2 qualifications": 20,
                "Level 3 qualifications": 20,
                "Level 4 qualifications or above": 50,
            },
        )
    )

    result = _qualification_percentages(df)

    assert result.columns[0] == "lsoa21"
    assert set(result.columns) == {"lsoa21", *(f"% {b}" for b in OUTPUT_BUCKETS)}
    row = result.filter(pl.col("lsoa21") == "E01000001").to_dicts()[0]
    assert row["% Degree or higher"] == 50.0
    assert row["% No qualifications"] == 10.0
    assert row["% Good GCSEs"] == 20.0
    assert row["% A-levels"] == 20.0
    # Percentages always sum to exactly 100 (largest-remainder rounding).
    assert round(sum(row[f"% {b}"] for b in OUTPUT_BUCKETS), 1) == 100.0


def test_qualification_colloquial_band_mapping():
    """ONS jargon bands fold into the colloquial buckets, not 'Level N' names."""
    df = pl.DataFrame(
        _long_rows(
            "E01000002",
            {
                "Level 1 and entry level qualifications": 40,
                "Apprenticeship": 60,
            },
        )
    )

    row = _qualification_percentages(df).to_dicts()[0]
    assert row["% Some GCSEs"] == 40.0
    assert row["% Apprenticeship"] == 60.0
    # No raw ONS "Level N" column names leak through.
    assert not any("Level" in c for c in row)


def test_qualification_percentages_independent_per_lsoa():
    df = pl.concat(
        [
            pl.DataFrame(_long_rows("E01000010", {"Level 4 qualifications or above": 100})),
            pl.DataFrame(_long_rows("E01000011", {"No qualifications": 100})),
        ]
    )

    result = _qualification_percentages(df).sort("lsoa21")
    assert result["% Degree or higher"].to_list() == [100.0, 0.0]
    assert result["% No qualifications"].to_list() == [0.0, 100.0]


def test_qualification_percentages_rejects_unexpected_band():
    rows = _long_rows("E01000003", {"Level 4 qualifications or above": 10})
    rows.append(
        {
            "GEOGRAPHY_CODE": "E01000003",
            "C2021_HIQUAL_8_NAME": "Level 5 brand new census band",
            "OBS_VALUE": 5,
        }
    )

    with pytest.raises(ValueError, match="do not match the expected"):
        _qualification_percentages(pl.DataFrame(rows))


def test_qualification_percentages_rejects_missing_band():
    rows = [
        r
        for r in _long_rows("E01000004", {"Level 4 qualifications or above": 10})
        if r["C2021_HIQUAL_8_NAME"] != "Apprenticeship"
    ]

    with pytest.raises(ValueError, match="missing"):
        _qualification_percentages(pl.DataFrame(rows))