#!/usr/bin/env python3
"""
deafrica_dem_to_minio.py
========================
One-off ETL: Digital Earth Africa DEM  →  Cloud-Optimized GeoTIFF  →  LUSPA MinIO.

This is Stage 1 of the LUPMIS2 GIS Analytical Tools concept: get a real raster
into the object store so the PWA can display it via MapView.addCOGLayer().

Designed to run inside the **DE Africa Sandbox** (which has the Open Data Cube
already configured).  It can also run locally, but note DE Africa's own docs:
locally, `dc.load` / `load_ard` need extra configuration (a database) — use the
Sandbox unless exports become routine.

    # 0. see which DEM products actually exist (do this first)
    python deafrica_dem_to_minio.py --list-products

    # 1. export + upload for a pilot district
    export MINIO_KEY=...        # never hardcode these
    export MINIO_SECRET=...
    python deafrica_dem_to_minio.py --district koforidua

    # or an explicit area / precise boundary
    python deafrica_dem_to_minio.py --bbox -0.35 6.00 -0.15 6.20 --name koforidua
    python deafrica_dem_to_minio.py --geojson district.geojson --name koforidua

Credentials come from the environment (MINIO_KEY / MINIO_SECRET) on purpose:
the existing PHP integration hardcodes them, and an earlier key pair ended up
committed to Gitea.  Keep them out of this file and out of git.
"""

import argparse
import json
import os
import sys
from datetime import datetime, timezone

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MINIO_ENDPOINT = "https://minioapi.lupmis4luspa.org"
MINIO_REGION   = "gh-greater-accra-luspa"
BUCKET         = "raster-objects"
KEY_PREFIX     = "dem"                      # objects land at dem/<name>_dem.tif

# Approximate bounding boxes (EPSG:4326: min_lon, min_lat, max_lon, max_lat).
# These are convenience defaults only — for production use --geojson with the
# authoritative district boundary from PostGIS so the clip matches the data.
DISTRICTS = {
    "koforidua": (-0.35, 6.00, -0.15, 6.20),   # New Juaben / Koforidua area
    "tamale":    (-1.10, 9.20, -0.60, 9.60),   # Tamale metropolitan area
}

# DE Africa DEM product. Verify against --list-products before a real run:
# DE Africa publishes SRTM ('dem_srtm') and its derivatives ('dem_srtm_deriv',
# the source of the slope WMS layer LUPMIS2 already uses).
DEFAULT_PRODUCT = "dem_srtm"

# EPSG:3857 matches the map view, so the browser does no reprojection.
# NOTE: for Stage 3 hydrology/slope, prefer a metric CRS (UTM 30N = EPSG:32630
# for western Ghana, 31N = EPSG:32631 for the east) — Web Mercator distorts
# distance with latitude and will bias slope and flow calculations.
DEFAULT_CRS = "EPSG:3857"
DEFAULT_RES = 30                            # metres — SRTM native resolution


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def list_products():
    """Print datacube products whose name mentions elevation/DEM/SRTM."""
    import datacube
    dc = datacube.Datacube(app="lupmis2_dem_discovery")
    products = dc.list_products()
    mask = products["name"].str.contains("dem|srtm|elev", case=False, na=False)
    hits = products[mask][["name", "description"]]
    if hits.empty:
        print("No DEM-like products found. All available products:\n")
        print(products[["name", "description"]].to_string())
    else:
        print("DEM-related products:\n")
        print(hits.to_string(index=False))


def bbox_from_geojson(path):
    """Bounding box (min_lon, min_lat, max_lon, max_lat) of a GeoJSON file."""
    with open(path) as fh:
        gj = json.load(fh)
    xs, ys = [], []

    def walk(coords):
        if isinstance(coords[0], (int, float)):
            xs.append(coords[0]); ys.append(coords[1]); return
        for c in coords:
            walk(c)

    feats = gj.get("features", [gj])
    for f in feats:
        geom = f.get("geometry", f)
        if geom and geom.get("coordinates"):
            walk(geom["coordinates"])
    if not xs:
        raise SystemExit(f"No coordinates found in {path}")
    return (min(xs), min(ys), max(xs), max(ys))


def load_dem(bbox, product, crs, res):
    """Load the DEM for a bounding box and return a 2-D DataArray."""
    import datacube
    dc = datacube.Datacube(app="lupmis2_dem_etl")
    min_lon, min_lat, max_lon, max_lat = bbox

    print(f"  loading '{product}' for bbox {bbox} at {res} m in {crs} …")
    ds = dc.load(
        product=product,
        x=(min_lon, max_lon),
        y=(min_lat, max_lat),
        output_crs=crs,
        resolution=(-res, res),
        # DEMs are static, so any single observation is the whole story.
        dask_chunks={},
    )
    if not ds.data_vars:
        raise SystemExit(
            f"'{product}' returned no data for this area. "
            f"Check the product name with --list-products and confirm coverage."
        )

    # Take the first band (SRTM DEMs expose a single elevation band) and drop
    # the time dimension if the product carries one.
    band = list(ds.data_vars)[0]
    da = ds[band]
    if "time" in da.dims:
        da = da.isel(time=0)
    print(f"  band '{band}', shape {tuple(da.shape)}")
    return da.compute() if hasattr(da, "compute") else da


def upload(local_path, key, endpoint, bucket, region):
    """Upload to MinIO. Credentials come from the environment."""
    import boto3
    from botocore.client import Config

    access = os.environ.get("MINIO_KEY")
    secret = os.environ.get("MINIO_SECRET")
    if not access or not secret:
        raise SystemExit(
            "Set MINIO_KEY and MINIO_SECRET in the environment.\n"
            "  export MINIO_KEY=...\n  export MINIO_SECRET=..."
        )

    s3 = boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access,
        aws_secret_access_key=secret,
        region_name=region,
        config=Config(s3={"addressing_style": "path"}),   # MinIO is path-style
    )
    size_mb = os.path.getsize(local_path) / 1e6
    print(f"  uploading {size_mb:.1f} MB → s3://{bucket}/{key} …")
    s3.upload_file(
        local_path, bucket, key,
        # Serve a real raster content-type; the PWA's COG pre-flight check
        # rejects text/html and reports the type it actually received.
        ExtraArgs={"ContentType": "image/tiff"},
    )
    return f"{endpoint}/{bucket}/{key}"


def verify(url):
    """Confirm the object is anonymously readable and supports range requests."""
    import urllib.request
    import urllib.error

    def probe(headers=None):
        req = urllib.request.Request(url, headers=headers or {})
        return urllib.request.urlopen(req, timeout=30)

    try:
        r = probe()
        ok_get = r.status == 200
        ctype = r.headers.get("content-type")
    except urllib.error.HTTPError as e:
        print(f"  ✗ anonymous GET failed: HTTP {e.code} — is the bucket policy public?")
        return False
    except Exception as e:                                   # noqa: BLE001
        print(f"  ✗ anonymous GET failed: {e}")
        return False

    try:
        r2 = probe({"Range": "bytes=0-9"})
        ok_range = r2.status == 206
    except urllib.error.HTTPError as e:
        ok_range = e.code == 206
    except Exception:                                        # noqa: BLE001
        ok_range = False

    print(f"  {'✓' if ok_get else '✗'} anonymous GET (content-type: {ctype})")
    print(f"  {'✓' if ok_range else '✗'} range requests (206 Partial Content)")
    return ok_get and ok_range


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--list-products", action="store_true",
                    help="list DEM-related datacube products and exit")
    ap.add_argument("--district", choices=sorted(DISTRICTS),
                    help="use a preset (approximate) district bounding box")
    ap.add_argument("--bbox", nargs=4, type=float,
                    metavar=("MIN_LON", "MIN_LAT", "MAX_LON", "MAX_LAT"))
    ap.add_argument("--geojson", help="clip to the bounding box of this GeoJSON (preferred)")
    ap.add_argument("--name", help="short name used in the object key")
    ap.add_argument("--product", default=DEFAULT_PRODUCT)
    ap.add_argument("--crs", default=DEFAULT_CRS)
    ap.add_argument("--res", type=float, default=DEFAULT_RES, help="resolution in metres")
    ap.add_argument("--outdir", default=".")
    ap.add_argument("--dry-run", action="store_true",
                    help="write the COG locally but do not upload")
    args = ap.parse_args()

    if args.list_products:
        list_products()
        return

    # ---- resolve the area of interest ----
    if args.geojson:
        bbox = bbox_from_geojson(args.geojson)
        name = args.name or os.path.splitext(os.path.basename(args.geojson))[0]
    elif args.district:
        bbox = DISTRICTS[args.district]
        name = args.name or args.district
        print("NOTE: preset bounding boxes are approximate. For production use "
              "--geojson with the authoritative district boundary.")
    elif args.bbox:
        bbox = tuple(args.bbox)
        name = args.name or "aoi"
    else:
        ap.error("choose an area: --district, --bbox or --geojson")

    name = name.lower().replace(" ", "_")
    fname = os.path.join(args.outdir, f"{name}_dem.tif")
    key = f"{KEY_PREFIX}/{name}_dem.tif"

    print(f"\nLUPMIS2 · DE Africa DEM → MinIO")
    print(f"  area '{name}'  bbox={bbox}")

    # ---- load + write COG ----
    from datacube.utils.cog import write_cog
    da = load_dem(bbox, args.product, args.crs, args.res)

    print(f"  writing COG → {fname}")
    out = write_cog(da, fname=fname, overwrite=True)
    if hasattr(out, "compute"):          # dask-backed writes are lazy
        out.compute()

    if args.dry_run:
        print(f"\nDry run complete: {fname} ({os.path.getsize(fname)/1e6:.1f} MB)")
        return

    # ---- upload + verify ----
    url = upload(fname, key, MINIO_ENDPOINT, BUCKET, MINIO_REGION)
    print(f"\n  object URL: {url}")
    ok = verify(url)

    manifest = {
        "name": f"{name.title()} DEM (SRTM 30 m, DE Africa)",
        "url": url,
        "key": key,
        "product": args.product,
        "bbox_4326": list(bbox),
        "crs": args.crs,
        "resolution_m": args.res,
        "created": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "verified": bool(ok),
    }
    mpath = os.path.join(args.outdir, f"{name}_dem.manifest.json")
    with open(mpath, "w") as fh:
        json.dump(manifest, fh, indent=2)

    print(f"  manifest  → {mpath}")
    if ok:
        print("\nReady. In LUPMIS2: Add External Layer → COG → paste the object URL above.")
    else:
        print("\nUploaded, but the public read/range check failed — the PWA will not "
              "be able to stream it until that is resolved.")
        sys.exit(1)


if __name__ == "__main__":
    main()
