#!/usr/bin/env python3
"""Download SRTM GL1 GeoTIFF tiles for a bounding box via OpenTopography."""

from __future__ import annotations

import argparse
import math
import sys
from pathlib import Path

import httpx

OPENTOPOGRAPHY_URL = "https://portal.opentopography.org/API/globaldem"


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)):
    ns = "N" if lat >= 0 else "S"
    for lng in range(math.floor(min_lng), math.ceil(max_lng)):
      ew = "E" if lng >= 0 else "W"
      tiles.append(f"{ns}{abs(lat):02d}{ew}{abs(lng):03d}")
  return tiles


def download_tile(tile: str, output_dir: Path, api_key: str | None) -> None:
  params = {
    "demtype": "SRTMGL1",
    "south": int(tile[1:3]) * (1 if tile[0] == "N" else -1),
    "north": int(tile[1:3]) * (1 if tile[0] == "N" else -1) + 1,
    "west": int(tile[4:7]) * (1 if tile[3] == "E" else -1),
    "east": int(tile[4:7]) * (1 if tile[3] == "E" else -1) + 1,
    "outputFormat": "GTiff",
  }
  if api_key:
    params["API_Key"] = api_key

  output_path = output_dir / f"{tile}.tif"
  if output_path.exists():
    print(f"skip {tile}: already exists")
    return

  print(f"downloading {tile}...")
  with httpx.stream("GET", OPENTOPOGRAPHY_URL, params=params, 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}")


def main() -> int:
  parser = argparse.ArgumentParser(description="Download SRTM GL1 GeoTIFF 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 names, e.g. N59E030")
  parser.add_argument("--output", default="data/srtm", help="Output directory")
  parser.add_argument("--api-key", default=None, help="OpenTopography API key")
  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:
    try:
      download_tile(tile, output_dir, args.api_key)
    except httpx.HTTPError as exc:
      print(f"failed {tile}: {exc}", file=sys.stderr)

  return 0


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