#!/usr/bin/env python3
"""Download RTS probability tiles from the WMTS service.

Specify either a bounding box (EPSG:3857) or explicit col/row ranges.

Examples
--------
# By bounding box:
python download_tiles.py --bbox -2382320 1900885 -1305316 2949040 --output my_tiles/

# By col/row range:
python download_tiles.py --col 5 5 --row 216 263 --output my_tiles/
"""

import argparse
import sys
from pathlib import Path

from owslib.wmts import WebMapTileService

WMTS_URL = "https://arcticdata.io/data/10.18739/A26970107/rts_probability/WMTSCapabilities.xml"
LAYER = "rts_probability"
TILEMATRIXSET = "WebMercatorQuad"
Z_LEVEL = "10"
FORMAT = "image/geotiff"

WEB_MERCATOR_MIN = -20037508.342789244
WEB_MERCATOR_MAX = 20037508.342789244
TILE_SPAN_Z10 = 39135.75848201


def bbox_to_tile_indices(bbox_minx, bbox_miny, bbox_maxx, bbox_maxy):
    bbox_minx = max(bbox_minx, WEB_MERCATOR_MIN)
    bbox_miny = max(bbox_miny, WEB_MERCATOR_MIN)
    bbox_maxx = min(bbox_maxx, WEB_MERCATOR_MAX)
    bbox_maxy = min(bbox_maxy, WEB_MERCATOR_MAX)

    col_min = max(0, int((bbox_minx - WEB_MERCATOR_MIN) / TILE_SPAN_Z10 + 1e-9))
    col_max = min(1023, int((bbox_maxx - WEB_MERCATOR_MIN) / TILE_SPAN_Z10 - 1e-9))
    row_min = max(0, int((WEB_MERCATOR_MAX - bbox_maxy) / TILE_SPAN_Z10 + 1e-9))
    row_max = min(1023, int((WEB_MERCATOR_MAX - bbox_miny) / TILE_SPAN_Z10 - 1e-9))

    return col_min, col_max, row_min, row_max


def download_tiles(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"Connecting to WMTS service...")

    try:
        wmts = WebMapTileService(WMTS_URL)
    except Exception as e:
        print(f"ERROR: Could not connect to WMTS service: {e}")
        return False

    if LAYER not in wmts.contents:
        print(f"ERROR: Layer '{LAYER}' not found. Available: {list(wmts.contents.keys())}")
        return False

    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"{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 RTS probability GeoTIFF tiles from the WMTS service.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

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

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

    args = parser.parse_args()

    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()

    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(col_min, col_max, row_min, row_max, args.output)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()