#!/usr/bin/env python3
"""
Download satellite tiles from Mapy.com REST API for Hradec Králové.
Zoom 18 ≈ 0.38 m/pixel at lat 50° — good for car detection.

API key: https://developer.mapy.com/  (zdarma po registraci)
"""

import math
import os
import time
import argparse
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
from PIL import Image

# Hradec Králové bounding box (min_lon, min_lat, max_lon, max_lat)
HK_BBOX = (15.7800, 50.1800, 15.8800, 50.2400)

# Mapy.com REST API v1 — mapset "aerial" = satelitní snímky
# Aerial: ČR až zoom 20, zbytek světa jen zoom 13
TILE_URL = "https://api.mapy.com/v1/maptiles/aerial/256/{z}/{x}/{y}?lang=cs&apikey={apikey}"
TILE_SIZE = 256
DEFAULT_ZOOM = 18
OUT_DIR = Path("tiles")


def lat_lon_to_tile(lat: float, lon: float, zoom: int) -> tuple[int, int]:
    """Convert lat/lon to tile x, y at given zoom (Slippy Map convention)."""
    n = 2 ** zoom
    x = int((lon + 180) / 360 * n)
    lat_rad = math.radians(lat)
    y = int((1 - math.log(math.tan(lat_rad) + 1 / math.cos(lat_rad)) / math.pi) / 2 * n)
    return x, y


def tile_to_lat_lon(x: int, y: int, zoom: int) -> tuple[float, float]:
    """Return NW corner (lat, lon) of tile."""
    n = 2 ** zoom
    lon = x / n * 360 - 180
    lat_rad = math.atan(math.sinh(math.pi * (1 - 2 * y / n)))
    lat = math.degrees(lat_rad)
    return lat, lon


def ground_resolution(lat: float, zoom: int) -> float:
    """Return meters per pixel at given lat and zoom."""
    return 156543.03 * math.cos(math.radians(lat)) / (2 ** zoom)


def download_tile(x: int, y: int, zoom: int, out_dir: Path, session: requests.Session, apikey: str) -> Path | None:
    path = out_dir / f"{zoom}_{x}_{y}.jpg"
    if path.exists():
        return path

    url = TILE_URL.format(z=zoom, x=x, y=y, apikey=apikey)
    try:
        r = session.get(url, timeout=15)
        r.raise_for_status()
        path.write_bytes(r.content)
        return path
    except Exception as e:
        print(f"  FAIL {x},{y}: {e}")
        return None


def stitch_tiles(tiles: dict[tuple, Path], x_range: range, y_range: range, zoom: int, out_path: Path):
    """Stitch downloaded tiles into one big image."""
    cols = len(x_range)
    rows = len(y_range)
    canvas = Image.new("RGB", (cols * TILE_SIZE, rows * TILE_SIZE))

    for j, y in enumerate(y_range):
        for i, x in enumerate(x_range):
            key = (x, y)
            if key in tiles and tiles[key]:
                try:
                    img = Image.open(tiles[key])
                    canvas.paste(img, (i * TILE_SIZE, j * TILE_SIZE))
                except Exception as e:
                    print(f"  Stitch error {x},{y}: {e}")

    canvas.save(out_path)
    print(f"Stitched image saved: {out_path}  ({canvas.width}x{canvas.height} px)")
    return canvas


def main():
    parser = argparse.ArgumentParser(description="Download Mapy.cz satellite tiles for Hradec Králové")
    parser.add_argument("--zoom", type=int, default=DEFAULT_ZOOM,
                        help=f"Zoom level (default: {DEFAULT_ZOOM})")
    parser.add_argument("--bbox", nargs=4, type=float,
                        metavar=("MIN_LON", "MIN_LAT", "MAX_LON", "MAX_LAT"),
                        default=HK_BBOX,
                        help="Bounding box (default: Hradec Králové)")
    parser.add_argument("--out", type=Path, default=OUT_DIR,
                        help="Output directory for tiles")
    parser.add_argument("--stitch", action="store_true",
                        help="Stitch all tiles into one image")
    parser.add_argument("--workers", type=int, default=8,
                        help="Parallel download workers")
    parser.add_argument("--delay", type=float, default=0.05,
                        help="Delay between requests (seconds)")
    parser.add_argument("--apikey", type=str, default=os.environ.get("MAPY_API_KEY", ""),
                        help="Mapy.com API key (nebo env MAPY_API_KEY)")
    args = parser.parse_args()

    if not args.apikey:
        parser.error("API klíč je povinný — použij --apikey nebo nastav env MAPY_API_KEY")

    min_lon, min_lat, max_lon, max_lat = args.bbox
    zoom = args.zoom

    res = ground_resolution((min_lat + max_lat) / 2, zoom)
    print(f"Zoom {zoom} → ~{res:.2f} m/pixel")

    # Tile range (NW corner has smaller y value)
    x_min, y_max = lat_lon_to_tile(min_lat, min_lon, zoom)
    x_max, y_min = lat_lon_to_tile(max_lat, max_lon, zoom)

    x_range = range(x_min, x_max + 1)
    y_range = range(y_min, y_max + 1)
    total = len(x_range) * len(y_range)

    print(f"Tiles: {len(x_range)} cols × {len(y_range)} rows = {total} total")
    print(f"Area: ~{len(x_range)*TILE_SIZE*res/1000:.1f} km × {len(y_range)*TILE_SIZE*res/1000:.1f} km")

    args.out.mkdir(parents=True, exist_ok=True)

    session = requests.Session()
    session.headers.update({
        "User-Agent": "Mozilla/5.0 (compatible; satellite-downloader/1.0)",
        "Referer": "https://mapy.cz/",
    })

    tiles: dict[tuple, Path | None] = {}
    done = 0

    jobs = [(x, y) for y in y_range for x in x_range]

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(download_tile, x, y, zoom, args.out, session, args.apikey): (x, y)
                   for x, y in jobs}
        for future in as_completed(futures):
            x, y = futures[future]
            tiles[(x, y)] = future.result()
            done += 1
            if done % 20 == 0 or done == total:
                print(f"  {done}/{total} tiles downloaded", end="\r")
            time.sleep(args.delay)

    ok = sum(1 for v in tiles.values() if v)
    print(f"\nDownloaded: {ok}/{total} tiles → {args.out}/")

    if args.stitch:
        stitch_out = args.out / f"hk_z{zoom}_stitched.jpg"
        stitch_tiles(tiles, x_range, y_range, zoom, stitch_out)


if __name__ == "__main__":
    main()