#! /usr/bin/env python3
# Copyright (C) 2016-2026 SpaceKnow, Inc.

"""Example usage of most important SpaceKnow APIs, see API Quickstart in
https://docs.spaceknow.com/.

usage:
    1. export SK_USER_EMAIL=your-user-email
    2. export SK_USER_PASSWORD=your-password
    3. ./spaceknow-api-demo.py
"""

import json
import os
from pathlib import Path
import time

import requests


SK_USER_EMAIL = os.environ.get("SK_USER_EMAIL")
SK_USER_PASSWORD = os.environ.get("SK_USER_PASSWORD")

# Authorization configuration
OAUTH_URL = "https://spaceknow.auth0.com/oauth/ro"
OAUTH_CLIENT_ID = "hmWJcfhRouDOaJK2L8asREMlMrv3jFE1"

# URLs of API endpoints
RAGNAR_SEARCH_URL = "https://api.spaceknow.com/imagery/search"
GET_STATUS_URL = "https://api.spaceknow.com/tasking/get-status"
KRAKEN_IMAGERY_GEOJSON_URL = "https://api.spaceknow.com/kraken/release"
KRAKEN_GRID_URL_PATTERN = (
    "https://api.spaceknow.com/kraken/grid/{map_id}/-/{z}/{x}/{y}/png"
)

# provider and dataset as documented on
# https://docs.spaceknow.com/api/ragnar.html
PROVIDER = "satellogic"
DATASET = "newsat_sr"
START_DATETIME = "2022-05-08 00:00:00"
END_DATETIME = "2022-05-08 23:59:59"
MAP_TYPE = "imagery"

# Description of area of interest. To display or generate new GeoJSON go to
# http://geojson.io/
# The example is a small rectangular polygon near Houston, Texas.
EXTENT_GEOJSON = {
    "type": "Polygon",
    "coordinates": [
        [
            [-96.21853637695312, 29.78368782378684],
            [-96.21853637695312, 29.792745768504258],
            [-96.22897338867188, 29.792745768504258],
            [-96.22897338867188, 29.78368782378684],
            [-96.21853637695312, 29.78368782378684],
        ]
    ],
}


def main() -> None:
    if not SK_USER_EMAIL or not SK_USER_PASSWORD:
        raise Exception(
            "Please export your Spaceknow user email and password to environment "
            "variables. Run: `export SK_USER_EMAIL=your-user-email; "
            "export SK_USER_PASSWORD=your-password`"
        )

    print("Getting authorization JWT...")
    token = get_auth_token()
    print(f"Auth token is: {token}")

    print("Searching available scenes...")
    scene_metadata = search_scenes(EXTENT_GEOJSON, PROVIDER, DATASET, token)
    metadata_str = json.dumps(scene_metadata[:3], indent=4, sort_keys=True)
    print(f"Found {len(scene_metadata)} scenes, first 3: \n{metadata_str}")

    first_scene = scene_metadata[0]

    kraken_response = release_kraken_map(first_scene, token)
    print(f"Release Kraken response: {kraken_response}")
    map_type_response = kraken_response[MAP_TYPE]
    tiles = map_type_response["tiles"]
    map_id = map_type_response["mapId"]

    # Notice there is no need for auth token for downloading Kraken tiles.
    # we can download coarse tiles directly mentioned in response:
    download_kraken_tiles(map_id, tiles)
    print(f"Written {len(tiles)} tiles mentioned in response.")

    # Or, we can download tiles at finer resolution by "generating" tiles at
    # a finer zoom level, up to max zoom:
    fine_tiles = generate_tiles_up_to_zoom(tiles, map_type_response["maxZoom"])
    download_kraken_tiles(map_id, fine_tiles)


def get_auth_token():
    """Request JWT token from Auth0 with your SpaceKnow account."""
    body = {
        "client_id": OAUTH_CLIENT_ID,
        "username": SK_USER_EMAIL,
        "password": SK_USER_PASSWORD,
        "connection": "Username-Password-Authentication",
        "grant_type": "password",
        "scope": "openid",
    }
    requests_response = requests.post(OAUTH_URL, json=body, timeout=(3.05, 180))
    requests_response.raise_for_status()
    return requests_response.json()["id_token"]


def search_scenes(extent, provider, dataset, token):
    """Search endpoint returns all scene metadata which intersect with requested extent
    and has intersection ratio greater or equal than given threshold.
    """
    body = {
        "datasets": [{"provider": provider, "dataset": dataset}],
        "startDatetime": START_DATETIME,
        "endDatetime": END_DATETIME,
        "extent": extent,
        "minIntersection": 0,
    }
    response = async_api_request(RAGNAR_SEARCH_URL, body, token)
    # TODO: Ragnar API search responses are paged. This loads the first page only.
    return response["results"]


def release_kraken_map(scene_metadata, token):
    """Release Kraken map and return response, mainly map ID. With ``allocate``
    the scene is bought, which usually costs credits.
    """
    scene_id = scene_metadata["sceneId"]
    print(f"Requesting Kraken imagery map for scene id {scene_id}...")
    body = {
        "mapTypes": [MAP_TYPE],
        "sceneIds": [scene_id],
        "extent": EXTENT_GEOJSON,
        "allocate": True,
    }
    return async_api_request(KRAKEN_IMAGERY_GEOJSON_URL, body, token)


def download_kraken_tiles(map_id, tiles):
    """Save Web Map Tiles of given map ID to disk."""
    for z, x, y in tiles:
        url = KRAKEN_GRID_URL_PATTERN.format(map_id=map_id, z=z, x=x, y=y)
        tile_response = requests.get(url, timeout=(3.05, 180))
        try:
            tile_response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            # 404 can happen when scene doesn't cover given tile or a portion
            # of image covering it is completely invalid, see
            # https://docs.spaceknow.com/api/mechanics.html#kraken-api-error-codes
            error_message = json.loads(e.response.content)
            if error_message["error"] == "TILE-DOES-NOT-EXIST":
                continue
            raise

        filename = f"{z}-{x}-{y}-truecolor.png"
        Path(filename).write_bytes(tile_response.content)
        print(f"Written {filename}.")


def generate_tiles_up_to_zoom(tiles, max_zoom):
    """Split tiles into their children so that their z is >= max_zoom."""
    for z, x, y in tiles:
        if z >= max_zoom:
            yield z, x, y
        else:
            children = [
                (z + 1, 2 * x, 2 * y),
                (z + 1, 2 * x, 2 * y + 1),
                (z + 1, 2 * x + 1, 2 * y),
                (z + 1, 2 * x + 1, 2 * y + 1),
            ]
            yield from generate_tiles_up_to_zoom(children, max_zoom)


def async_api_request(url, body, token):
    """Initiate pipeline, check pipeline status until it finishes and
    then retrieve results.
    """
    initiate_response = api_request(url + "/initiate", body, token)
    pipeline_id = initiate_response["pipelineId"]
    print(f"Created pipeline {pipeline_id}, polling its status...")

    polling_body = {"pipelineId": pipeline_id}
    exponential_backoff_delay_s = 5  # start with 5 second intervals
    while True:
        time.sleep(exponential_backoff_delay_s)
        get_status_response = api_request(GET_STATUS_URL, polling_body, token)
        status = get_status_response["status"]
        print(f"Slept {exponential_backoff_delay_s} s, pipeline status now: {status}")

        if status not in {"NEW", "PROCESSING"}:
            break

        exponential_backoff_delay_s *= 1.4  # increase delay between polls

    return api_request(url + "/retrieve", polling_body, token)


def api_request(url, body, token):
    """Make HTTP POST request to SpaceKnow API and return response as dict."""
    headers = {"Authorization": f"Bearer {token}"}
    requests_response = requests.post(
        url, json=body, headers=headers, timeout=(3.05, 180)
    )

    # TODO: we should retry with exponential back-off for HTTP 500 class errors
    requests_response.raise_for_status()
    return requests_response.json()


if __name__ == "__main__":
    # SpaceKnow APIs return interesting payload even in error HTTP responses:
    try:
        main()
    except requests.HTTPError as e:
        print(f"HTTP error response payload: {e.response.text}")
        raise