"""Download one ECMWF SEAS5 initialisation from the CDS API.

LT=7 is requested only as the 00 UTC boundary after the end of LT=6.  The
monthly forecasting products themselves use LT=1..6.
"""
from __future__ import annotations

import argparse
from datetime import date
from pathlib import Path

from .config import PROJECT_ROOT


DATASET = "seasonal-original-single-levels"
VARIABLES = [
    "total_precipitation", "2m_temperature", "sea_surface_temperature",
    "soil_temperature_level_1", "surface_latent_heat_flux",
    "surface_sensible_heat_flux", "top_net_solar_radiation",
    "mean_sea_level_pressure", "top_net_thermal_radiation",
    "volumetric_soil_water_layer_1",
]


def _next_month(year: int, month: int, offset: int) -> tuple[int, int]:
    month_index = year * 12 + (month - 1) + offset
    return month_index // 12, month_index % 12 + 1


def lead_hours_through_lt7_boundary(year: int, month: int) -> list[str]:
    """Six-hourly lead hours from init through 00 UTC on first day of LT=7."""
    end_year, end_month = _next_month(year, month, 7)
    hours = (date(end_year, end_month, 1) - date(year, month, 1)).days * 24
    return [str(hour) for hour in range(0, hours + 1, 6)]


def retrieve(init: str, output: Path, variables: list[str]) -> None:
    try:
        import cdsapi
    except ImportError as exc:
        raise SystemExit("cdsapi is required; install it in the active environment.") from exc
    year, month = map(int, init.split("-"))
    if not 1 <= month <= 12:
        raise ValueError("--init must be formatted YYYY-MM with a valid month")
    output.parent.mkdir(parents=True, exist_ok=True)
    request = {
        "originating_centre": "ecmwf",
        "system": "51",
        "variable": variables,
        "year": str(year),
        "month": f"{month:02d}",
        "day": "01",
        "leadtime_hour": lead_hours_through_lt7_boundary(year, month),
        "data_format": "netcdf"
    }
    cdsapi.Client().retrieve(DATASET, request, str(output))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--init", required=True, help="Initial month: YYYY-MM")
    parser.add_argument("--predictands-only", action="store_true",
                        help="Download only total precipitation and 2m temperature")
    parser.add_argument("--output", type=Path, help="Optional target NetCDF path")
    args = parser.parse_args()
    target = args.output or PROJECT_ROOT / "data/ecmwf/raw" / f"ecmwf_seas5_{args.init}.nc"
    variables = VARIABLES[:2] if args.predictands_only else VARIABLES
    retrieve(args.init, target, variables)


if __name__ == "__main__":
    main()