#!/usr/bin/env python3
"""
Download gridMET daily NetCDF files into year-based folders.

Default layout:
data/gridmet/1991/sph.nc
data/gridmet/1991/rmax.nc
data/gridmet/1991/rmin.nc

gridMET source:
https://www.northwestknowledge.net/metdata/data/{variable}_{year}.nc
"""

from __future__ import annotations

import argparse
import sys
import tempfile
import urllib.error
import urllib.request
from pathlib import Path


DEFAULT_BASE_URL = "https://www.northwestknowledge.net/metdata/data"
DEFAULT_VARIABLES = ("sph", "rmax", "rmin")


def _parse_years(text: str) -> list[int]:
    years: set[int] = set()
    for part in text.split(","):
        token = part.strip()
        if not token:
            continue
        if "-" in token:
            start_text, end_text = token.split("-", 1)
            start = int(start_text)
            end = int(end_text)
            if start > end:
                raise argparse.ArgumentTypeError(f"Invalid descending year range: {token}")
            years.update(range(start, end + 1))
        else:
            years.add(int(token))
    if not years:
        raise argparse.ArgumentTypeError("At least one year is required.")
    return sorted(years)


def _parse_variables(text: str) -> list[str]:
    variables = [part.strip() for part in text.split(",") if part.strip()]
    if not variables:
        raise argparse.ArgumentTypeError("At least one variable is required.")
    return variables


def _download_file(url: str, destination: Path, overwrite: bool) -> bool:
    if destination.exists() and not overwrite:
        print(f"skip existing {destination}")
        return False

    destination.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        prefix=destination.stem + ".",
        suffix=".download",
        dir=destination.parent,
        delete=False,
    ) as handle:
        temp_path = Path(handle.name)

    try:
        print(f"download {url}")
        with urllib.request.urlopen(url) as response, temp_path.open("wb") as output:
            total_text = response.headers.get("Content-Length")
            total = int(total_text) if total_text and total_text.isdigit() else None
            copied = 0
            while True:
                chunk = response.read(1024 * 1024)
                if not chunk:
                    break
                output.write(chunk)
                copied += len(chunk)
                if total:
                    percent = copied / total * 100
                    print(f"\r  {copied / 1024 / 1024:,.1f} MiB / {total / 1024 / 1024:,.1f} MiB ({percent:5.1f}%)", end="")
            if total:
                print()

        temp_path.replace(destination)
        print(f"saved {destination}")
        return True
    except (urllib.error.URLError, urllib.error.HTTPError, OSError) as exc:
        temp_path.unlink(missing_ok=True)
        print(f"failed {url}: {exc}", file=sys.stderr)
        raise


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Download gridMET NetCDF files into data/gridmet/<year>/<variable>.nc."
    )
    parser.add_argument(
        "--years",
        type=_parse_years,
        default=_parse_years("1991-2020"),
        help="Comma-separated years/ranges. Default: 1991-2020.",
    )
    parser.add_argument(
        "--variables",
        type=_parse_variables,
        default=list(DEFAULT_VARIABLES),
        help="Comma-separated gridMET variables. Default: sph,rmax,rmin.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/gridmet"),
        help="Root folder for downloaded files. Default: data/gridmet.",
    )
    parser.add_argument(
        "--base-url",
        default=DEFAULT_BASE_URL,
        help=f"Base gridMET download URL. Default: {DEFAULT_BASE_URL}.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Redownload files that already exist.",
    )
    args = parser.parse_args()

    downloaded = 0
    for year in args.years:
        for variable in args.variables:
            url = f"{args.base_url.rstrip('/')}/{variable}_{year}.nc"
            destination = args.output_dir / str(year) / f"{variable}.nc"
            if _download_file(url, destination, args.overwrite):
                downloaded += 1

    print(f"done; downloaded {downloaded} file(s)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())