#!/usr/bin/env python3
"""Download Pan-Arctic Biomass Mapping tiles from the WMTS service.

The service (see WMTScopy.xml) publishes many layers, all sharing a single
custom tile matrix set ("Custom3571TMS") in EPSG:3571 with only one zoom
level (matrix "0", a 23x23 grid). Specify the layer to download, and either
a bounding box (EPSG:3571) or explicit col/row ranges.

Examples
--------
# List available layers:
python download-tiles.py --list-layers

# By bounding box:
python download-tiles.py --layer AGB_2020_Plant_p500 --bbox -2000000 -2000000 0 0 --output my_tiles/

# By col/row range:
python download-tiles.py --layer AGB_2020_Plant_p500 --col 5 8 --row 10 12 --output my_tiles/
"""

import argparse
import sys
from pathlib import Path

from owslib.wmts import WebMapTileService

WMTS_URL = "https://arcticdata.io/data/10.18739/A2GF0MZ6W/WMTSCapabilities.xml"
TILEMATRIXSET = "Custom3571TMS"
Z_LEVEL = "0"
FORMAT = "image/geotiff"

# From the Custom3571TMS TileMatrixSet definition (EPSG:3571, matrix "0")
TMS_MIN = -4889340.0
TMS_MAX = 4889340.0
MATRIX_WIDTH = 23
MATRIX_HEIGHT = 23
TILE_SPAN_X = (TMS_MAX - TMS_MIN) / MATRIX_WIDTH
TILE_SPAN_Y = (TMS_MAX - TMS_MIN) / MATRIX_HEIGHT


def bbox_to_tile_indices(bbox_minx, bbox_miny, bbox_maxx, bbox_maxy):
    bbox_minx = max(bbox_minx, TMS_MIN)
    bbox_miny = max(bbox_miny, TMS_MIN)
    bbox_maxx = min(bbox_maxx, TMS_MAX)
    bbox_maxy = min(bbox_maxy, TMS_MAX)

    col_min = max(0, int((bbox_minx - TMS_MIN) / TILE_SPAN_X + 1e-9))
    col_max = min(MATRIX_WIDTH - 1, int((bbox_maxx - TMS_MIN) / TILE_SPAN_X - 1e-9))
    row_min = max(0, int((TMS_MAX - bbox_maxy) / TILE_SPAN_Y + 1e-9))
    row_max = min(MATRIX_HEIGHT - 1, int((TMS_MAX - bbox_miny) / TILE_SPAN_Y - 1e-9))

    return col_min, col_max, row_min, row_max


def download_tiles(wmts, layer, col_min, col_max, row_min, row_max, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    total = (col_max - col_min + 1) * (row_max - row_min + 1)
    print(f"Downloading up to {total} tiles: Col [{col_min}–{col_max}], Row [{row_min}–{row_max}]")
    print(f"Output: {output_dir.resolve()}\n")

    downloaded = skipped = failed = 0

    for col in range(col_min, col_max + 1):
        for row in range(row_min, row_max + 1):
            try:
                tile = wmts.gettile(
                    layer=layer,
                    tilematrixset=TILEMATRIXSET,
                    tilematrix=Z_LEVEL,
                    column=col,
                    row=row,
                    format=FORMAT,
                )
                out_path = output_dir / f"{layer}_{col}_{row}.tif"
                out_path.write_bytes(tile.read())
                downloaded += 1
                if downloaded % 20 == 0:
                    print(f"  {downloaded} downloaded so far...")

            except Exception as e:
                msg = str(e)
                if "404" in msg or "Not Found" in msg:
                    # All-NoData tiles are intentionally absent
                    skipped += 1
                else:
                    failed += 1
                    print(f"  FAILED col={col} row={row}: {e}")

    print(f"\nDone: {downloaded} downloaded, {skipped} skipped (NoData), {failed} failed")
    return failed == 0


def parse_args():
    parser = argparse.ArgumentParser(
        description="Download Pan-Arctic Biomass Mapping GeoTIFF tiles from the WMTS service.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        "--layer",
        metavar="IDENTIFIER",
        help="Layer identifier to download (see --list-layers)",
    )
    parser.add_argument(
        "--list-layers",
        action="store_true",
        help="List available layer identifiers and exit",
    )

    area = parser.add_mutually_exclusive_group(required=False)
    area.add_argument(
        "--bbox",
        nargs=4,
        type=float,
        metavar=("MINX", "MINY", "MAXX", "MAXY"),
        help="Bounding box in EPSG:3571 metres (west south east north)",
    )
    area.add_argument(
        "--col",
        nargs=2,
        type=int,
        metavar=("COL_MIN", "COL_MAX"),
        help=f"Column range (0–{MATRIX_WIDTH - 1}, west to east)",
    )

    parser.add_argument(
        "--row",
        nargs=2,
        type=int,
        metavar=("ROW_MIN", "ROW_MAX"),
        help="Row range (0–%d, north to south). Required when using --col." % (MATRIX_HEIGHT - 1),
    )
    parser.add_argument(
        "--output",
        default="tiles_output",
        metavar="DIR",
        help="Output directory (default: tiles_output/)",
    )

    args = parser.parse_args()

    if args.list_layers:
        return args

    if not args.layer:
        parser.error("--layer is required")
    if args.bbox is None and args.col is None:
        parser.error("one of --bbox or --col is required")
    if args.col is not None and args.row is None:
        parser.error("--row is required when using --col")

    return args


def main():
    args = parse_args()

    print("Connecting to WMTS service...")
    try:
        wmts = WebMapTileService(WMTS_URL)
    except Exception as e:
        print(f"ERROR: Could not connect to WMTS service: {e}")
        sys.exit(1)

    if args.list_layers:
        print("Available layers:")
        for identifier in wmts.contents:
            print(f"  {identifier}")
        sys.exit(0)

    if args.layer not in wmts.contents:
        print(f"ERROR: Layer '{args.layer}' not found. Available: {list(wmts.contents.keys())}")
        sys.exit(1)

    if args.bbox:
        minx, miny, maxx, maxy = args.bbox
        col_min, col_max, row_min, row_max = bbox_to_tile_indices(minx, miny, maxx, maxy)
        print(f"Bounding box ({minx}, {miny}, {maxx}, {maxy})")
        print(f"→ Tile range: Col [{col_min}–{col_max}], Row [{row_min}–{row_max}]")
    else:
        col_min, col_max = args.col
        row_min, row_max = args.row

    if col_min > col_max or row_min > row_max:
        print("ERROR: Empty tile range — check that your bounding box or indices are valid.")
        sys.exit(1)

    success = download_tiles(wmts, args.layer, col_min, col_max, row_min, row_max, args.output)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()