# snap_request.py
"""Submit one catalogue L1C product and a complete graph to NSIS custom_snap.

With no action option, only check public catalogue metadata. Set the documented
NSIS_S3_* variables and ACCESS_TOKEN before using --submit or --resume.
Use --check-access to check the ordering API with ACCESS_TOKEN without creating
an order or accessing S3 storage.
"""
import argparse
import copy
import json
import os
from pathlib import Path
import time
from uuid import uuid4
from urllib.parse import urlparse

import requests

CATALOGUE = 'https://catalogue.nsiscloud.polsa.gov.pl/odata/v1'
ORDERING = 'https://odp.nsiscloud.polsa.gov.pl/odata/v1'
PRODUCT = 'S2A_MSIL1C_20251230T101451_N0511_R022_T33UVV_20251230T103257.SAFE'
GRAPH = Path(__file__).with_name('nsis_l1c_subset.xml')
OUTPUT_NAME = 'nsis_l1c_subset.tif'


def check_product(product_name):
    if '_MSIL1C_' not in product_name or not product_name.endswith('.SAFE'):
        raise ValueError('This graph example requires a Sentinel-2 L1C SAFE product.')
    response = requests.get(CATALOGUE + '/Products', params={
        '$filter': "Name eq '" + product_name.replace("'", "''") + "'",
        '$expand': 'Locations',
    }, timeout=60)
    response.raise_for_status()
    products = response.json()['value']
    if len(products) != 1 or products[0].get('Online') is not True:
        raise RuntimeError('The exact L1C product is not uniquely available online. Select and inspect a compatible L1C input first.')
    return products[0]


def build_order(product_name, endpoint, access_key, secret_key, bucket, run_id):
    products_in_graph = {'infile': {'source_type': 'catalogue', 'source': product_name}}
    output_filenames = {'outfile': OUTPUT_NAME}
    graph_key = f'graphs/{run_id}/{GRAPH.name}'
    output_prefix = f'snap-results/{run_id}'
    workflow_options = [
        {'Name': 's3_endpoint_url', 'Value': endpoint},
        {'Name': 's3_access_key', 'Value': access_key},
        {'Name': 's3_secret_key', 'Value': secret_key},
        {'Name': 'graph_path', 'Value': f's3://{bucket}/{graph_key}'},
        {'Name': 'products_in_graph', 'Value': products_in_graph},
        {'Name': 'output_filenames', 'Value': output_filenames},
        {'Name': 'output_storage', 'Value': 'PRIVATE'},
        {'Name': 'output_dir', 'Value': f's3://{bucket}/{output_prefix}'},
    ]
    order_body = {
        'WorkflowName': 'custom_snap',
        'Name': f'NSIS L1C subset {run_id}',
        'InputProductReference': {'Reference': ''},
        'WorkflowOptions': workflow_options,
    }
    return order_body, graph_key, f'{output_prefix}/{OUTPUT_NAME}'


def redacted_order(order_body):
    visible = copy.deepcopy(order_body)
    for option in visible['WorkflowOptions']:
        if option['Name'] in ('s3_access_key', 's3_secret_key'):
            option['Value'] = '[REDACTED]'
    return visible


def api_request(method, path, **kwargs):
    # Read ACCESS_TOKEN each time so a newly obtained token can be used on restart.
    token = os.environ.get('ACCESS_TOKEN')
    if not token:
        raise RuntimeError('Set ACCESS_TOKEN using the NSIS authentication instructions.')
    response = requests.request(method, ORDERING + path,
                                headers={'Authorization': f'Bearer {token}'},
                                timeout=90, **kwargs)
    if response.status_code in (401, 403):
        if method == 'GET' and path == '/Workflows':
            next_step = ('No order was submitted by this invocation. Sign in again in this '
                         'terminal, then run python3 snap_request.py --check-access.')
        elif method == 'GET' and path.startswith('/BatchOrder('):
            next_step = ('Sign in again in this terminal, then use --resume with the saved '
                         'run JSON to retrieve this order. Do not use --submit.')
        else:
            next_step = ('Check your orders in the NSIS portal before submitting again. '
                         'Use --resume if an order was created and its run JSON was saved.')
        raise RuntimeError(
            f'NSIS ordering API returned HTTP {response.status_code} for {method} {path}. '
            'The access token may have expired or access may have been denied. ' + next_step)
    if response.status_code == 504:
        if method == 'GET' and path == '/Workflows':
            next_step = ('No order was submitted by this invocation. Wait a minute, then '
                         'run python3 snap_request.py --check-access again.')
        elif method == 'GET' and path.startswith('/BatchOrder('):
            next_step = ('The order already exists. Wait a minute, then use --resume '
                         'with the saved run JSON. Do not use --submit.')
        else:
            next_step = ('The submission outcome is unknown. Check your orders in the '
                         'NSIS portal before submitting again.')
        raise RuntimeError(f'NSIS gateway timed out (HTTP 504) for {method} {path}. ' + next_step)
    response.raise_for_status()
    return response.json()


def statuses(payload):
    value = payload.get('value', payload)
    records = value if isinstance(value, list) else [value]
    if not records or any(not isinstance(x, dict) or 'Status' not in x for x in records):
        raise RuntimeError('Unexpected order status response; inspect it in the NSIS portal before continuing.')
    return [str(x['Status']).lower() for x in records]


def check_workflow():
    workflows = api_request('GET', '/Workflows', params={
        '$expand': 'WorkflowOptions', '$filter': "Name eq 'custom_snap'",
    })
    if len(workflows.get('value', [])) != 1:
        raise RuntimeError('The API did not return exactly one custom_snap workflow. '
                           'No order was submitted by this invocation.')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    action = parser.add_mutually_exclusive_group()
    action.add_argument('--submit', action='store_true', help='Upload the graph and submit a processing order')
    action.add_argument('--check-access', action='store_true', help='Check authenticated workflow access without creating an order or accessing S3')
    action.add_argument('--resume', type=Path, help='Resume status/output checks from a run JSON; never submit another order')
    parser.add_argument('--product', default=PRODUCT)
    args = parser.parse_args()
    if args.check_access:
        print('Checking ordering API access: GET /Workflows (request timeout: 90 seconds)...', flush=True)
        check_workflow()
        print('Ordering API access check passed: custom_snap is listed. No order was created. '
              'S3 access and order submission have not been tested.')
        return
    if args.submit and not GRAPH.is_file():
        raise FileNotFoundError(f'Missing SNAP graph: {GRAPH}. Save the XML as '
                                'nsis_l1c_subset.xml alongside snap_request.py. '
                                'No order was submitted by this invocation.')
    if not args.resume:
        product = check_product(args.product)
        print(json.dumps({key: product[key] for key in ('Id', 'Name', 'Online')}, indent=2))
        if not args.submit:
            print('Public metadata check only. Use --submit after preparing credentials and reviewing the graph.')
            return

    import boto3
    endpoint = os.environ['NSIS_S3_ENDPOINT']
    if urlparse(endpoint).scheme != 'https':
        raise ValueError('NSIS_S3_ENDPOINT must be your HTTPS private-storage endpoint.')
    bucket = os.environ['NSIS_S3_BUCKET']
    access_key = os.environ['NSIS_S3_ACCESS_KEY']
    secret_key = os.environ['NSIS_S3_SECRET_KEY']
    s3 = boto3.client('s3', endpoint_url=endpoint,
                      aws_access_key_id=access_key, aws_secret_access_key=secret_key)
    if args.resume:
        run = json.loads(args.resume.read_text())
        if run['endpoint'] != endpoint or run['bucket'] != bucket:
            raise ValueError('Resume with the same private-storage endpoint and bucket.')
    else:
        check_workflow()
        run_id = uuid4().hex
        order_body, graph_key, output_key = build_order(args.product, endpoint, access_key, secret_key, bucket, run_id)
        print(json.dumps(redacted_order(order_body), indent=2))
        s3.upload_file(str(GRAPH), bucket, graph_key)
        # Submit once: a network timeout after this POST needs an order-list check.
        created = api_request('POST', '/ProductionOrder/OData.CSC.Order', json=order_body)
        order_id = created['value']['Id']
        run = {'order_id': order_id, 'endpoint': endpoint, 'bucket': bucket,
               'output_key': output_key, 'product': args.product}
        run_file = Path(f'snap-run-{run_id}.json')
        run_file.write_text(json.dumps(run, indent=2) + '\n')
        print(f'Order created: {order_id}. Resume file: {run_file}')

    # Stop after one hour; the saved order can be checked again with --resume.
    for _ in range(120):
        payload = api_request('GET', f"/BatchOrder({run['order_id']})")
        states = statuses(payload)
        print('Order status:', ', '.join(states))
        if all(state in ('completed', 'done') for state in states):
            break
        if any(state in ('failed', 'error', 'cancelled', 'canceled') for state in states):
            raise RuntimeError('Processing failed. Inspect StatusMessage in the order response or portal.')
        time.sleep(30)
    else:
        raise TimeoutError('Order has not finished. Use --resume later; do not submit it again.')
    output = Path(f"{run['order_id']}-{OUTPUT_NAME}")
    if output.exists():
        raise FileExistsError(output)
    s3.download_file(bucket, run['output_key'], str(output))
    import rasterio
    import numpy as np
    with rasterio.open(output) as dataset:
        if dataset.count != 2 or (dataset.height, dataset.width) != (512, 512) or dataset.crs is None:
            raise RuntimeError('Unexpected GeoTIFF shape, bands or CRS; review the graph and service output.')
        pixels = dataset.read(masked=True).astype(float)
        if not np.any(np.isfinite(pixels.filled(np.nan)) & (pixels.filled(0) != 0)):
            raise RuntimeError('The selected subset has no usable pixels. Inspect the scene and change the graph region.')
    print('Readable two-band GeoTIFF:', output)


if __name__ == '__main__':
    main()