"""Submit an image and task in one request, then retrieve its result.

Requires: pip install requests
Set P5D_API_BASE_URL, P5D_API_TOKEN and P5D_REQUEST_ID before running.
Use P5D_JOB_ID to resume polling an already submitted job without creating another.
"""
import os
import json
import sys
import time
from urllib.parse import quote

import requests

api_base = os.environ["P5D_API_BASE_URL"].rstrip("/")
api_token = os.environ["P5D_API_TOKEN"]
maximum_credits = int(os.environ.get("P5D_MAXIMUM_CREDITS", "20"))
stage_auth = None
if os.environ.get("P5D_STAGE_USERNAME") and os.environ.get("P5D_STAGE_PASSWORD"):
    stage_auth = (os.environ["P5D_STAGE_USERNAME"], os.environ["P5D_STAGE_PASSWORD"])

headers = {"Accept": "application/json, application/problem+json"}
if stage_auth:
    headers["as"] = api_token  # Keep staging Basic Auth separate from the API credential.
else:
    headers["Authorization"] = "Bearer " + api_token


def api(method, path, **kwargs):
    response = requests.request(
        method, api_base + path, headers=headers,
        timeout=120, allow_redirects=False, auth=stage_auth, **kwargs
    )
    response.raise_for_status()
    if 300 <= response.status_code < 400:
        raise RuntimeError("Unexpected API redirect; check the API base URL")
    data = response.json()
    if isinstance(data.get("error"), int) and data["error"] != 0:
        raise RuntimeError(data)
    return response


job_id = os.environ.get("P5D_JOB_ID")
if not job_id:
    request_id = os.environ["P5D_REQUEST_ID"]
    source_filename = sys.argv[1] if len(sys.argv) > 1 else "room.jpg"
    task = {
        "operation": os.environ.get("P5D_OPERATION", "render-image"),
        "request_id": request_id,
        "instruction": "Natural daylight and neutral materials",
        "count": 1,
        "maximum_credits": maximum_credits,
    }
    print("Keep this request ID before submission:", request_id, flush=True)
    try:
        with open(source_filename, "rb") as source:
            response = api("POST", "/generations", files={"source": source},
                           data={"request": json.dumps(task)})
        job_id = response.json()["jobId"]
    except (requests.Timeout, requests.ConnectionError) as error:
        raise RuntimeError(
            "Submission outcome is unknown. Retry the same request ID, image bytes and task; "
            "do not assign a new request ID. Uploaded inputs and accepted work are reused."
        ) from error
    except requests.HTTPError as error:
        # A 409 with Location can identify work accepted before an upload receipt expired.
        location = error.response.headers.get("Location", "")
        if error.response.status_code != 409 or not location.startswith("/api/public/generations/"):
            raise
        job_id = location.rsplit("/", 1)[-1]
    print("Keep this job ID:", job_id, flush=True)

print("Polling:", job_id, flush=True)
deadline = time.monotonic() + 600
while time.monotonic() < deadline:
    response = api("GET", "/generations/" + quote(job_id, safe=""))
    job = response.json()
    if job["status"] == "succeeded":
        for result in job.get("results", []):
            print("Output:", result.get("src"))
        print("Billing:", api("GET", "/generations/" + quote(job_id, safe="") + "/billing").json())
        break
    if job["status"] == "failed":
        raise RuntimeError(job.get("error", "Generation failed"))
    retry_after = response.headers.get("Retry-After", "2")
    delay = max(1, int(retry_after)) if retry_after.isdigit() else 2
    time.sleep(min(delay, max(0, deadline - time.monotonic())))
else:
    raise TimeoutError("Still running. Resume with P5D_JOB_ID=" + job_id)