"""ONS (Office for National Statistics) UK population extractor.

Fetches mid-year population estimates by Local Authority from the ONS Dataset API.
Downloads a CSV (~68MB) containing population by LAD × year × sex × age, then
filters to sex='all' for the most recent year and sums across ages to get
total population per LAD.

No authentication required.

Landing: {LANDING_DIR}/ons_uk/{year}/{month}/lad_population.json.gz
Output:  {"rows": [{"lad_code": "E08000003", "lad_name": "Manchester",
                    "population": 553230, "ref_year": 2022,
                    "country_code": "GB"}], "count": N}
"""

import csv
import io
import json
import sqlite3
from pathlib import Path

import niquests

from ._shared import HTTP_TIMEOUT_SECONDS, run_extractor, setup_logging
from .utils import get_last_cursor, landing_path, write_gzip_atomic

logger = setup_logging("padelnomics.extract.ons_uk")

EXTRACTOR_NAME = "ons_uk"

ONS_BASE = "https://api.beta.ons.gov.uk/v1"
DATASET_ID = "mid-year-pop-est"
# Most recent edition with England & Wales LAD data
EDITION = "mid-2022-england-wales"

MIN_POPULATION = 50_000

# CSV columns: v4_0, calendar-years, Time, administrative-geography, Geography, sex, Sex, single-year-of-age, Age
COL_VALUE = "v4_0"
COL_YEAR = "calendar-years"
COL_LAD_CODE = "administrative-geography"
COL_LAD_NAME = "Geography"
COL_SEX = "sex"


def _get_csv_url(session: niquests.Session) -> tuple[str, int]:
    """Fetch the latest CSV download URL from the ONS Dataset API.

    Returns (url, ref_year).
    """
    versions_url = f"{ONS_BASE}/datasets/{DATASET_ID}/editions/{EDITION}/versions"
    resp = session.get(versions_url, timeout=HTTP_TIMEOUT_SECONDS)
    resp.raise_for_status()
    versions = resp.json().get("items", [])
    assert versions, f"No versions found for {DATASET_ID}/{EDITION}"

    # Sort by version number (highest = latest) and pick the latest
    latest = max(versions, key=lambda v: v.get("version", 0))
    csv_url = latest["downloads"]["csv"]["href"]

    # Extract year from edition name (mid-2022-england-wales → 2022)
    year_part = EDITION.split("-")[1]
    ref_year = int(year_part)

    return csv_url, ref_year


def _aggregate_csv(content: bytes, ref_year: int) -> list[dict]:
    """Parse the ONS CSV and aggregate total population per LAD for the target year.

    The CSV has one row per (year × LAD × sex × age); we filter to sex='all'
    and the target year, then sum the value column per LAD to get total population.
    """
    reader = csv.DictReader(io.StringIO(content.decode("utf-8")))
    totals: dict[str, dict] = {}
    target_year = str(ref_year)

    for row in reader:
        if row.get(COL_SEX, "").lower() != "all":
            continue
        if row.get(COL_YEAR, "") != target_year:
            continue
        lad_code = row.get(COL_LAD_CODE, "").strip()
        lad_name = row.get(COL_LAD_NAME, "").strip()
        if not lad_code or not lad_name:
            continue
        # Skip aggregate geographies (country/region level codes start differently from LADs)
        # LAD codes: E06*, E07*, E08*, E09*, W06*, S12*, N09*
        if not any(lad_code.startswith(p) for p in ("E0", "W0", "S1", "N0")):
            continue
        try:
            value = int(row.get(COL_VALUE, "0").replace(",", ""))
        except ValueError:
            continue
        if lad_code not in totals:
            totals[lad_code] = {"lad_code": lad_code, "lad_name": lad_name, "population": 0}
        totals[lad_code]["population"] += value

    return list(totals.values())


def extract(
    landing_dir: Path,
    year_month: str,
    conn: sqlite3.Connection,
    session: niquests.Session,
) -> dict:
    """Fetch ONS LAD population. Skips if already run this month."""
    last_cursor = get_last_cursor(conn, EXTRACTOR_NAME)
    if last_cursor == year_month:
        logger.info("already have data for %s — skipping", year_month)
        return {"files_written": 0, "files_skipped": 1, "bytes_written": 0}

    year, month = year_month.split("/")

    logger.info("fetching ONS CSV download URL for %s/%s", DATASET_ID, EDITION)
    csv_url, ref_year = _get_csv_url(session)

    logger.info("GET %s (ref_year=%d)", csv_url, ref_year)
    resp = session.get(csv_url, timeout=HTTP_TIMEOUT_SECONDS * 10)
    resp.raise_for_status()
    assert len(resp.content) > 1_000_000, (
        f"ONS CSV too small ({len(resp.content)} bytes) — download may have failed"
    )

    logger.info("aggregating %d bytes CSV", len(resp.content))
    lad_rows = _aggregate_csv(resp.content, ref_year)
    assert len(lad_rows) > 100, f"Expected >100 LADs from CSV, got {len(lad_rows)}"

    filtered = [
        {
            "lad_code": r["lad_code"],
            "lad_name": r["lad_name"],
            "population": r["population"],
            "ref_year": ref_year,
            "country_code": "GB",
        }
        for r in lad_rows
        if r["population"] >= MIN_POPULATION
    ]

    assert len(filtered) > 50, f"Expected >50 UK LADs ≥50K pop, got {len(filtered)}"
    logger.info("parsed %d UK LADs with population ≥%d", len(filtered), MIN_POPULATION)

    dest_dir = landing_path(landing_dir, "ons_uk", year, month)
    dest = dest_dir / "lad_population.json.gz"
    payload = json.dumps({"rows": filtered, "count": len(filtered)}).encode()
    bytes_written = write_gzip_atomic(dest, payload)
    logger.info("written %s bytes compressed", f"{bytes_written:,}")

    return {
        "files_written": 1,
        "files_skipped": 0,
        "bytes_written": bytes_written,
        "cursor_value": year_month,
    }


def main() -> None:
    run_extractor(EXTRACTOR_NAME, extract)


if __name__ == "__main__":
    main()
