#!/usr/bin/env python3
"""Download ETH Global Canopy Height 2020 tiles for a bounding box."""

from __future__ import annotations

import argparse
import math
import sys
from pathlib import Path

import httpx

# Public tile base used by the ETH Global Canopy Height dataset portal.
CANOPY_BASE_URL = "https://storage.googleapis.com/earthenginepartners-hansen/GCH2020"


def tile_names_for_bbox(min_lat: float, min_lng: float, max_lat: float, max_lng: float) -> list[str]:
  tiles: list[str] = []
  for lat in range(math.floor(min_lat), math.ceil(max_lat)):
    for lng in range(math.floor(min_lng), math.ceil(max_lng)):
      tiles.append(f"GCH2020_E{lng:03d}N{lat:02d}.tif")
  return tiles


def download_tile(tile: str, output_dir: Path) -> None:
  output_path = output_dir / tile
  if output_path.exists():
    print(f"skip {tile}: already exists")
    return

  url = f"{CANOPY_BASE_URL}/{tile}"
  print(f"downloading {tile}...")
  try:
    with httpx.stream("GET", url, timeout=300.0, follow_redirects=True) as response:
      response.raise_for_status()
      with output_path.open("wb") as file:
        for chunk in response.iter_bytes():
          file.write(chunk)
    print(f"saved {output_path}")
  except httpx.HTTPError as exc:
    print(f"failed {tile}: {exc}", file=sys.stderr)


def main() -> int:
  parser = argparse.ArgumentParser(description="Download ETH Global Canopy Height tiles")
  parser.add_argument("--bbox", nargs=4, type=float, metavar=("MIN_LAT", "MIN_LNG", "MAX_LAT", "MAX_LNG"))
  parser.add_argument("--tiles", nargs="*", help="Explicit tile filenames")
  parser.add_argument("--output", default="data/canopy", help="Output directory")
  args = parser.parse_args()

  output_dir = Path(args.output)
  output_dir.mkdir(parents=True, exist_ok=True)

  if args.tiles:
    tiles = args.tiles
  elif args.bbox:
    tiles = tile_names_for_bbox(*args.bbox)
  else:
    parser.error("provide --bbox or --tiles")

  for tile in tiles:
    download_tile(tile, output_dir)

  return 0


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