from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path

from countryatlas.connectors.owid import OWIDConnector
from countryatlas.models import IndicatorSourceSpec, RawPayload


def _conn() -> OWIDConnector:
    c = OWIDConnector.__new__(OWIDConnector)
    c._cache, c._frames = {}, {}
    import threading

    c._lock = threading.Lock()
    return c


def _raw(body: bytes, dataset: str, code: str) -> RawPayload:
    return RawPayload(connector="owid", dataset=dataset, code=code, url="http://test", retrieved_at=datetime.now(UTC),
                      status_code=200, content_type="text/csv", body=body)


def test_co2_subset_normalizes_per_column(fixtures: Path) -> None:
    body = (fixtures / "owid_co2_subset.csv").read_bytes()
    raw = _raw(body, "co2", "owid-co2-data")
    conn = _conn()
    rows = conn.normalize(raw, IndicatorSourceSpec(indicator_id="co2-per-capita", connector="owid", dataset="co2", code="co2_per_capita"))
    isos = {r.country_id for r in rows}
    assert isos == {"CAN", "FRA", "NGA", "KWT"}  # World (OWID_WRL) and Asia (no code) dropped
    assert all(r.unit == "tonnes per person" for r in rows)
    assert all(r.frequency == "A" and r.period.month == 1 for r in rows)
    # second spec on the same payload reuses the parsed frame
    rows2 = conn.normalize(raw, IndicatorSourceSpec(indicator_id="co2-emissions", connector="owid", dataset="co2", code="co2"))
    assert len(rows2) >= len(rows) - 5
    assert len(conn._frames) == 1


def test_grapher_projection_column_becomes_forecast(fixtures: Path) -> None:
    body = (fixtures / "owid_grapher_median-age.csv").read_bytes()
    conn = _conn()
    rows = conn.normalize(_raw(body, "grapher", "median-age"),
                          IndicatorSourceSpec(indicator_id="median-age", connector="owid", dataset="grapher", code="median-age"))
    can = {r.year: r for r in rows if r.country_id == "CAN"}
    assert can[2020].is_forecast is False
    assert any(r.is_forecast for r in can.values())
    assert max(can) <= datetime.now(UTC).year + 6
    assert "WORLD" not in {r.country_id for r in rows} and "OWID_WRL" not in {r.country_id for r in rows}