"""Download Census 2021 usual-resident counts per unit postcode.

ONS Census 2021 publishes a *direct* headcount of usual residents for every
unit postcode in England and Wales (table P001, disaggregated by sex) as a bulk
CSV on the NOMIS webservice. We sum the two sex rows per postcode to get the
total usual-resident population and emit one row per unit postcode.

This is a display-only side table (shown in the right-hand area pane); it is NOT
merged into the property/postcode feature frames and is never a filterable
attribute. The Rust server loads the parquet directly via --population-path.

Source: NOMIS bulk product "Postcode resident and household estimates, England
        and Wales: Census 2021" (table P001), as at census day 21 March 2021.
        https://www.nomisweb.co.uk/sources/census_2021_pc
License: Open Government Licence v3.0

Caveats (Census 2021 disclosure control): counts carry targeted record swapping
and cell-key perturbation, and only postcodes with at least one usual resident
are present, so vacant/non-residential postcodes are absent from the output.
"""

import argparse
import tempfile
import time
from pathlib import Path

import polars as pl

from pipeline.utils import download

# P001 = "Number of usual residents by postcode by sex", split alphabetically by
# postcode into four CSVs (each: Postcode, Sex Code, Sex Label, Count). The
# per-file minimums are ~95% of the published record counts and let us retry
# until a download is complete (see _fetch_csv).
BASE = "https://www.nomisweb.co.uk/output/census/2021"
P001_FILES = [
    ("pcd_p001_a_d.csv", 720_000),
    ("pcd_p001_e_l.csv", 540_000),
    ("pcd_p001_m_r.csv", 630_000),
    ("pcd_p001_s_z.csv", 750_000),
]

# P001 covers England & Wales (~1.37M postcodes). A materially smaller result
# means a split file was truncated or silently 404'd.
MIN_EXPECTED_POSTCODES = 1_000_000


def _canonical_postcode(col: pl.Expr) -> pl.Expr:
    """Canonical spaced unit postcode (e.g. "AL1 1AG"), matching the Rust
    server's `normalize_postcode`: uppercase, strip non-alphanumerics, then
    insert a single space before the final three (inward-code) characters."""
    cleaned = col.cast(pl.String).str.to_uppercase().str.replace_all(r"[^A-Z0-9]", "")
    return (
        cleaned.str.slice(0, cleaned.str.len_chars() - 3)
        + pl.lit(" ")
        + cleaned.str.slice(-3)
    )


def _fetch_csv(
    url: str, dest: Path, min_rows: int, *, max_attempts: int = 20
) -> pl.DataFrame:
    """Download a CSV, defending against silent truncation.

    The NOMIS bulk files are served with chunked transfer encoding and no
    Content-Length, so a dropped connection ends the stream without raising —
    yielding a short file. We retry until the parsed row count reaches the
    file's known approximate size, keeping the largest parse seen.
    """
    best: pl.DataFrame | None = None
    for attempt in range(1, max_attempts + 1):
        try:
            download(url, dest)
            df = pl.read_csv(dest)
        except Exception as exc:  # noqa: BLE001 — retry any transport/parse error
            print(f"  {url} attempt {attempt}: error {exc}")
            time.sleep(3)
            continue
        rows = df.height
        if best is None or rows > best.height:
            best = df
        complete = rows >= min_rows
        print(
            f"  {url} attempt {attempt}: {rows} rows "
            f"({'complete' if complete else f'< {min_rows}, retrying'})"
        )
        if complete:
            return df
        time.sleep(2)
    raise RuntimeError(
        f"Failed to fully download {url} after {max_attempts} attempts "
        f"(best {best.height if best is not None else 0} rows, need >= {min_rows})"
    )


def download_and_convert(output_path: Path) -> None:
    frames: list[pl.DataFrame] = []
    with tempfile.TemporaryDirectory() as tmp:
        for name, min_rows in P001_FILES:
            url = f"{BASE}/{name}"
            dest = Path(tmp) / name
            print(f"Downloading {url} ...")
            frames.append(_fetch_csv(url, dest, min_rows))

    df = pl.concat(frames)
    print(f"Total P001 rows (postcode x sex): {df.height}")

    if "Count" not in df.columns or "Postcode" not in df.columns:
        raise ValueError(
            f"Unexpected P001 columns {df.columns}; expected 'Postcode' and 'Count'."
        )

    result = (
        df.with_columns(_canonical_postcode(pl.col("Postcode")).alias("postcode"))
        .filter(pl.col("postcode").str.len_chars() >= 5)
        .group_by("postcode")
        .agg(pl.col("Count").cast(pl.Int64).sum().alias("population"))
        .filter(pl.col("population") > 0)
        .with_columns(pl.col("population").cast(pl.UInt32))
        .sort("postcode")
    )

    print(f"Unit postcodes with population: {result.height}")
    if result.height < MIN_EXPECTED_POSTCODES:
        raise ValueError(
            f"Only {result.height} postcodes (expected >= {MIN_EXPECTED_POSTCODES}); "
            "a NOMIS P001 split file was likely truncated or unavailable."
        )
    total = result["population"].sum()
    print(f"Total usual residents (England & Wales): {total:,}")
    print(
        f"Per-postcode population range: {result['population'].min()} - "
        f"{result['population'].max()}"
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.write_parquet(output_path, compression="zstd")
    print(f"Saved to {output_path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download Census 2021 usual-resident counts per unit postcode"
    )
    parser.add_argument(
        "--output", type=Path, required=True, help="Output parquet file path"
    )
    args = parser.parse_args()
    download_and_convert(args.output)


if __name__ == "__main__":
    main()