from __future__ import annotations

from datetime import UTC, date, datetime

import polars as pl

from countryatlas.models import NormalizedObservation
from countryatlas.pipeline.staging import rows_to_frame
from countryatlas.pipeline.validate import validate_frame
from countryatlas.registry import indicators_by_id


def _rows(ind: str, unit: str, series: dict[str, list[float]], start: int = 2000) -> pl.DataFrame:
    now = datetime.now(UTC)
    out = []
    for c, vals in series.items():
        for k, v in enumerate(vals):
            out.append(NormalizedObservation(country_id=c, indicator_id=ind, period=date(start + k, 1, 1), year=start + k,
                                             frequency="A", value=v, unit=unit, source_id="worldbank", source_dataset="WDI",
                                             source_series_code="X", retrieved_at=now))
    return rows_to_frame(out)


def test_bounds_quarantine_and_jump_warning() -> None:
    ind = indicators_by_id()["life-expectancy"]  # bounds roughly [20, 100]
    df = _rows("life-expectancy", ind.unit, {"CAN": [78, 78.5, 79, 79.4, 79.9, 80.2, 80.6, 81.0, 81.3, 250.0],
                                               "FRA": [78, 78.4, 78.9, 79.3, 79.7, 80.1, 80.5, 60.0, 81.2, 81.5]}, 2010)
    gv = validate_frame(df, ind, now=datetime(2021, 6, 1, tzinfo=UTC))
    st = dict(zip(gv.frame["country_id"].to_list(), gv.frame["status"].to_list(), strict=True))  # last per country wins
    assert gv.frame.filter((pl.col("country_id") == "CAN") & (pl.col("year") == 2019))["status"][0] == "quarantined"
    assert gv.frame.filter((pl.col("country_id") == "FRA") & (pl.col("year") == 2017))["status"][0] == "warning"
    assert gv.n_quarantined == 1 and gv.n_warning >= 1
    assert not gv.quarantine_dataset
    assert {i.code for i in gv.issues} >= {"out_of_bounds", "extreme_jump"}
    assert st  # sanity


def test_unit_mismatch_and_partial_download_quarantine_dataset() -> None:
    ind = indicators_by_id()["gdp-per-capita"]
    df = _rows("gdp-per-capita", "wrong unit", {"CAN": [1.0, 2.0, 3.0]})
    gv = validate_frame(df, ind)
    assert gv.quarantine_dataset and any(i.code == "unit_mismatch" for i in gv.issues)
    df2 = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]})
    gv2 = validate_frame(df2, ind, previous_rows=100)
    assert gv2.quarantine_dataset and any(i.code == "partial_download" for i in gv2.issues)
    gv3 = validate_frame(df2, ind, previous_rows=5)
    assert not gv3.quarantine_dataset


def test_stale_flag_on_latest_row_only() -> None:
    ind = indicators_by_id()["gdp-per-capita"]
    df = _rows("gdp-per-capita", ind.unit, {"CAN": [100.0, 101.0, 102.0, 103.0, 104.0]}, 2010)  # latest 2014
    gv = validate_frame(df, ind, now=datetime(2026, 9, 1, tzinfo=UTC))
    statuses = gv.frame.sort("period")["status"].to_list()
    assert statuses[-1] == "stale" and statuses[:-1] == ["imported"] * 4