"""Download NaPTAN data and extract railway/metro station POIs."""

import argparse
import io
import urllib.request
from pathlib import Path

import polars as pl

NAPTAN_CSV_URL = "https://naptan.api.dft.gov.uk/v1/access-nodes?dataFormat=csv"


STOP_TYPES = {
    "AIR": "Airport",
    "FTD": "Ferry",
    "RSE": "Rail station",
    "BCT": "Bus stop",
    "BCE": "Bus station",
    "TXR": "Taxi rank",
    "TMU": "Metro or Tram stop",
}


def download_naptan(output: Path) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)

    print(f"Downloading NaPTAN data from {NAPTAN_CSV_URL}")
    with urllib.request.urlopen(NAPTAN_CSV_URL) as resp:
        raw = resp.read()

    print(f"Downloaded {len(raw) / (1024 * 1024):.1f} MB")

    df = (
        pl.read_csv(io.BytesIO(raw), infer_schema_length=0)
        .with_columns(
            pl.col("Latitude").cast(pl.Float64, strict=False),
            pl.col("Longitude").cast(pl.Float64, strict=False),
        )
        .drop_nulls(subset=["Latitude", "Longitude"])
        .filter(pl.col("StopType").is_in(list(STOP_TYPES.keys())))
        .select(
            pl.col("ATCOCode").alias("id"),
            pl.col("CommonName").alias("name"),
            pl.col("StopType").replace(STOP_TYPES).alias("category"),
            pl.col("Latitude").alias("lat"),
            pl.col("Longitude").alias("lng"),
        )
    )

    df.write_parquet(output)
    size_mb = output.stat().st_size / (1024 * 1024)
    print(f"Wrote {output} ({size_mb:.1f} MB, {len(df):,} stations)")

    counts = df.group_by("category").len().sort("len", descending=True)
    for row in counts.iter_rows(named=True):
        print(f"  {row['category']}: {row['len']:,}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Download NaPTAN station data")
    parser.add_argument(
        "--output", type=Path, required=True, help="Output parquet file path"
    )
    args = parser.parse_args()
    download_naptan(args.output)


if __name__ == "__main__":
    main()