"""Submit one SimAtomic simulation and download its results.

Requires: pip install requests
Set SIMATOMIC_API_KEY, then run: python api_workflow_example.py input.zip
Running this example uploads the archive and submits a compute job.
"""
import os
from pathlib import Path
import sys
import time
from urllib.parse import urljoin, urlparse

import requests

BASE = "https://app.simatomic.com/api/api_handler/"


def main():
    archive = Path(sys.argv[1])
    headers = {"X-API-Key": os.environ["SIMATOMIC_API_KEY"]}

    def post(route, payload):
        response = requests.post(BASE + route, json=payload, headers=headers,
                                 timeout=60, allow_redirects=False)
        if route == "poll_job" and response.status_code == 404:
            return {"job_status": "queued"}
        response.raise_for_status()
        if response.is_redirect:
            raise RuntimeError("Unexpected API redirect")
        return response.json()

    # Short workflow test: choose production length for your scientific question.
    params = {"mode": "simulation", "md_steps": 1000}
    with archive.open("rb") as source:
        upload = post("get_presigned_url", {"key": archive.name})
        upload_url = upload["presigned_url"]
        if urlparse(upload_url).scheme != "https":
            raise RuntimeError("Expected an HTTPS upload URL")
        response = requests.put(upload_url, data=source, timeout=300,
                                allow_redirects=False)
        response.raise_for_status()
        if response.is_redirect:
            raise RuntimeError("Unexpected upload redirect")

    job_id = upload["job_id"]
    post("queue_job", {**params, "key": archive.name, "job_id": job_id})
    print("Submitted job:", job_id, flush=True)
    post("start_remote_server", {})

    deadline = time.monotonic() + 24 * 60 * 60
    while time.monotonic() < deadline:
        status = post("poll_job", {"job_id": job_id})
        state = status.get("job_status", "unknown")
        print("Status:", state, flush=True)
        if state == "success":
            break
        if state in {"failed", "error", "aborted", "cancelled", "canceled"}:
            raise RuntimeError(f"Job {job_id} ended: {state}")
        time.sleep(60)
    else:
        raise TimeoutError(f"Polling timed out; job {job_id} may still be running")

    # Send the API key only to the SimAtomic API. Follow the result redirect
    # without that header, using the signed HTTPS result URL.
    response = requests.get(BASE + "download", params={"job_id": job_id},
                            headers=headers, timeout=60, stream=True,
                            allow_redirects=False)
    for _ in range(5):
        if not response.is_redirect:
            break
        location = urljoin(response.url, response.headers["Location"])
        response.close()
        if urlparse(location).scheme != "https":
            raise RuntimeError("Expected an HTTPS result URL")
        response = requests.get(location, stream=True, timeout=300,
                                allow_redirects=False)
    with response:
        response.raise_for_status()
        if response.is_redirect:
            raise RuntimeError("Too many result redirects")
        target = Path("results") / f"{job_id}.zip"
        target.parent.mkdir(exist_ok=True)
        temporary = target.with_suffix(".zip.part")
        with temporary.open("wb") as output:
            for block in response.iter_content(1024 * 1024):
                output.write(block)
        temporary.replace(target)
    print("Saved:", target)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit("Usage: python api_workflow_example.py input.zip")
    main()