#!/usr/bin/env python3
"""Downloads the 1 x 1 degree cells of Earth's surface at 10 m that scenes name, and
makes the renderer's fine regions from them.

    scripts/fetch-earth-regions.py [--scene FILE] [--box S W N E] [CELL ...]

Cells come from each scene's `[appearance] earth_regions` boxes, from each --box
(south, west, north, east in degrees; west east of east crosses the 180 degree
meridian), and from names such as N29E032 (29 to 30 N, 32 to 33 E). Only those cells
are fetched; nothing else is, and the renderer never fetches anything itself.

Source: ESA WorldCover 10 m 2021 v200, Sentinel-2 RGBNIR composite: the median of the
year's cloud-free Sentinel-2 L2A surface reflectance in bands B04, B03, B02 and B08,
as 12000 x 12000 GeoTIFFs of 16-bit values (1e-4 reflectance per unit, 0 where there
is no data), one per cell with land in it, from the public bucket
s3://esa-worldcover-s2 (Zanaga et al. 2022, doi:10.5281/zenodo.7254221). License:
CC BY 4.0; credit "ESA WorldCover project 2021 / Contains modified Copernicus Sentinel
data (2021) processed by ESA WorldCover consortium".

Each file (about 650 MB) goes to data/earth/regions/source/ and is checked against
the bucket's ETag, the MD5 of its 8 MiB parts. A cell the bucket has no file for is
all water, and gets an empty data/earth/regions/<cell>.none so the renderer does not
ask for it again. Then `resolved earth-regions` makes data/earth/regions/<cell>.tiles
(194 MB of BC6H, about 10 s a cell; it builds the CLI first), and remakes cells made
before whose neighbors have arrived since, so their borders join up and their holes
are filled alike.

Needs only the Python standard library.
"""

import argparse
import hashlib
import math
import subprocess
import sys
import tomllib
import urllib.error
import urllib.request
from pathlib import Path

BUCKET = "https://esa-worldcover-s2.s3.eu-central-1.amazonaws.com/rgbnir/2021"
PART = 8 * 1024 * 1024  # the bucket's multipart upload part size
ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "data" / "earth" / "regions"
SOURCE = OUT / "source"


def name(south, west):
    """The cell's name, as N29E032 (as Cell::name in earth.rs)."""
    ns = "S" if south < 0 else "N"
    ew = "W" if west < 0 else "E"
    return f"{ns}{abs(south):02}{ew}{abs(west):03}"


def parse(cell):
    """South and west of a name such as N29E032."""
    ok = len(cell) == 7 and cell[0] in "NS" and cell[3] in "EW"
    ok = ok and cell[1:3].isdigit() and cell[4:].isdigit()
    if not ok:
        sys.exit(f"{cell!r} does not name a cell, as N29E032 does")
    south = int(cell[1:3]) * (-1 if cell[0] == "S" else 1)
    west = int(cell[4:]) * (-1 if cell[3] == "W" else 1)
    if not (-90 <= south < 90 and -180 <= west < 180):
        sys.exit(f"{cell} is not on Earth")
    return south, west


def covering(south, west, north, east):
    """The cells over a box, as Cell::covering in earth.rs."""
    valid = -90 <= south < north <= 90 and -180 <= west <= 180
    valid = valid and -180 <= east <= 180 and west != east
    if not valid:
        sys.exit(f"{south}, {west} to {north}, {east} is not a box of latitude and longitude")
    first, last = math.floor(west), math.ceil(east)
    if last <= first:
        last += 360
    cells = set()
    for s in range(max(math.floor(south), -90), min(math.ceil(north), 90)):
        for w in range(first, last):
            cells.add((s, (w + 180) % 360 - 180))
    return cells


def file_name(cell):
    return f"ESA_WorldCover_10m_2021_v200_{cell}_S2RGBNIR.tif"


def url(cell):
    return f"{BUCKET}/{cell[:3]}/{file_name(cell)}"


def etag(cell):
    """The bucket's ETag for the cell, or None if it has no file for it."""
    request = urllib.request.Request(url(cell), method="HEAD")
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            return response.headers["ETag"].strip('"')
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return None
        raise


def local_etag(path, parts):
    """The ETag S3 gives a file uploaded in `parts` parts of PART bytes, or whole."""
    if parts == 0:
        return hashlib.md5(path.read_bytes()).hexdigest()
    digests = []
    with open(path, "rb") as f:
        while chunk := f.read(PART):
            digests.append(hashlib.md5(chunk).digest())
    return f"{hashlib.md5(b''.join(digests)).hexdigest()}-{len(digests)}"


def matches(path, tag):
    parts = int(tag.split("-")[1]) if "-" in tag else 0
    return local_etag(path, parts) == tag


def fetch(cell):
    """Downloads the cell's file if it is not here and checked; False if the bucket
    has none."""
    tag = etag(cell)
    if tag is None:
        (OUT / f"{cell}.none").touch()
        print(f"{cell}: all water, no file")
        return False
    path = SOURCE / file_name(cell)
    if path.is_file() and matches(path, tag):
        return True
    print(f"{cell}: fetching")
    part = path.with_suffix(".part")
    with urllib.request.urlopen(url(cell), timeout=60) as response, open(part, "wb") as f:
        while chunk := response.read(1 << 20):
            f.write(chunk)
    if not matches(part, tag):
        part.unlink()
        sys.exit(f"{cell}: download does not match the bucket's ETag {tag}")
    part.rename(path)
    return True


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--scene", action="append", default=[], type=Path)
    parser.add_argument("--box", action="append", default=[], nargs=4, type=float)
    parser.add_argument("cells", nargs="*")
    args = parser.parse_args()

    wanted = {parse(c) for c in args.cells}
    for box in args.box:
        wanted |= covering(*box)
    for scene in args.scene:
        with open(scene, "rb") as f:
            boxes = tomllib.load(f).get("appearance", {}).get("earth_regions", [])
        if not boxes:
            print(f"{scene} names no earth_regions")
        for box in boxes:
            wanted |= covering(*box)
    if not wanted:
        sys.exit("no cells named: give --scene, --box or cell names")

    OUT.mkdir(parents=True, exist_ok=True)
    SOURCE.mkdir(parents=True, exist_ok=True)
    names = sorted(name(s, w) for s, w in wanted)
    fetched = [c for c in names if not (OUT / f"{c}.none").is_file() and fetch(c)]
    # Cells made before are passed too: those next to new ones are remade.
    made = {p.stem for p in OUT.glob("*.tiles")}
    build = sorted(set(fetched) | made)
    if build:
        subprocess.run(
            ["cargo", "run", "--quiet", "--release", "-p", "resolved-cli", "--",
             "earth-regions", "--source", str(SOURCE), "--out", str(OUT), *build],
            cwd=ROOT, check=True,
        )
    print(f"{len(fetched)} cells in {OUT}")


if __name__ == "__main__":
    main()