import polars as pl
import pytest

from pipeline.download.tenure import OUTPUT_BUCKETS, TENURE_MAP, _tenure_percentages


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

    Every one of the 8 leaf categories must be present in the download (NOMIS
    emits a 0-count row when an LSOA has none), so categories not given default
    to 0 to mirror that.
    """
    return [
        {
            "GEOGRAPHY_CODE": geo,
            "C2021_TENURE_9_NAME": label,
            "OBS_VALUE": counts.get(label, 0),
        }
        for label in TENURE_MAP
    ]


def test_tenure_percentages_keyed_by_lsoa_with_three_buckets():
    df = pl.DataFrame(
        _long_rows(
            "E01000001",
            {
                "Owned: Owns outright": 40,
                "Owned: Owns with a mortgage or loan": 15,
                "Shared ownership: Shared ownership": 5,
                "Social rented: Rents from council or Local Authority": 15,
                "Social rented: Other social rented": 5,
                "Private rented: Private landlord or letting agency": 18,
                "Private rented: Other private rented": 1,
                "Lives rent free": 1,
            },
        )
    )

    result = _tenure_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]
    # Owner occupied = outright + mortgage + shared ownership = 60.
    assert row["% Owner occupied"] == 60.0
    # Social rent = council + other social = 20.
    assert row["% Social rent"] == 20.0
    # Private rent = private landlord + other private + lives rent free = 20.
    assert row["% Private rent"] == 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_tenure_shared_ownership_rolls_into_owner_occupied():
    """Shared ownership is part-owned, so it counts as owner-occupied, not rent."""
    df = pl.DataFrame(_long_rows("E01000002", {"Shared ownership: Shared ownership": 100}))

    row = _tenure_percentages(df).to_dicts()[0]

    assert row["% Owner occupied"] == 100.0
    assert row["% Social rent"] == 0.0
    assert row["% Private rent"] == 0.0


def test_tenure_lives_rent_free_rolls_into_private_rent():
    """'Lives rent free' folds into private rent (ONS 'private rented or rent free')."""
    df = pl.DataFrame(_long_rows("E01000003", {"Lives rent free": 100}))

    row = _tenure_percentages(df).to_dicts()[0]

    assert row["% Private rent"] == 100.0
    assert row["% Owner occupied"] == 0.0
    assert row["% Social rent"] == 0.0


def test_tenure_percentages_independent_per_lsoa():
    """Two LSOAs get independent profiles: the LSOA granularity is the point."""
    df = pl.concat(
        [
            pl.DataFrame(_long_rows("E01000010", {"Owned: Owns outright": 100})),
            pl.DataFrame(
                _long_rows(
                    "E01000011",
                    {"Social rented: Rents from council or Local Authority": 100},
                )
            ),
        ]
    )

    result = _tenure_percentages(df).sort("lsoa21")

    assert result["% Owner occupied"].to_list() == [100.0, 0.0]
    assert result["% Social rent"].to_list() == [0.0, 100.0]


def test_tenure_percentages_rejects_unexpected_category():
    rows = _long_rows("E01000004", {"Owned: Owns outright": 10})
    rows.append(
        {
            "GEOGRAPHY_CODE": "E01000004",
            "C2021_TENURE_9_NAME": "A Brand New Census Tenure Category",
            "OBS_VALUE": 5,
        }
    )

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


def test_tenure_percentages_rejects_missing_category():
    # Drop one leaf entirely: its households would vanish from the denominator.
    rows = [
        r
        for r in _long_rows("E01000005", {"Owned: Owns outright": 10})
        if r["C2021_TENURE_9_NAME"] != "Lives rent free"
    ]

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