"""Upload, add a periodic bond, optimise and download using the StructSub API.

Use an existing server; the default URL assumes a configured development checkout. Install the client with `python -m pip install httpx`.
Run this script in a folder where you want to save the two output files.
Local use needs no API key. STRUCTSUB_TOKEN is an optional hosted sign-in token,
not a personal API key; hosted sign-in tokens expire.
"""

import json
import os
import time
from pathlib import Path

import httpx

BASE_URL = os.environ.get("STRUCTSUB_URL", "http://127.0.0.1:8000").rstrip("/")
TOKEN = os.environ.get("STRUCTSUB_TOKEN")

# Two H atoms whose closest connection crosses the a boundary.
# The 1.4 Å separation is deliberately stretched, so we add the bond ourselves.
SAMPLE = '''2
Lattice="12 0 0 0 12 0 0 0 12" Properties=species:S:1:pos:R:3 pbc="T T T"
H 0.2 0 0
H 10.8 0 0
'''


def finish(client, response, timeout=300):
    """Return an inline result, or wait for a queued job with a finite timeout."""
    response.raise_for_status()
    result = response.json()
    if not result.get("queued"):
        return result
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        status = client.get(f"/jobs/{result['job']}")
        status.raise_for_status()
        job = status.json()
        if job["state"] == "done":
            return job["result"]
        if job["state"] in {"failed", "cancelled"}:
            raise RuntimeError(job.get("error") or f"Job {job['state']}")
        time.sleep(0.5)
    raise TimeoutError("Stopped waiting; the server job may still be running.")


def main():
    headers = {"Authorization": f"Bearer {TOKEN}"} if TOKEN else {}
    # One client retains trial cookies if using a hosted deployment.
    with httpx.Client(base_url=BASE_URL, headers=headers, timeout=120) as client:
        uploaded = finish(client, client.post(
            "/sessions", files={"file": ("stretched-hydrogen.extxyz", SAMPLE.encode())},
        ))
        session = uploaded["session"]
        print(f"Session: {session}")
        if uploaded.get("warnings"):
            print("Connectivity warnings:", uploaded["warnings"])

        finish(client, client.post(f"/sessions/{session}/bonds", json={
            "i": 0, "j": 1, "offset": [-1, 0, 0],
        }))
        relaxed = finish(client, client.post(f"/sessions/{session}/relax", json={
            "engine": "uff-periodic", "scope": "all", "max_steps": 200,
            "optimise_cell": False, "background": True,
        }))
        print(relaxed["report"]["summary"])
        print("Converged:", relaxed["report"]["converged"])

        exported = client.get(f"/sessions/{session}/download", params={"format": "cif"})
        exported.raise_for_status()
        Path("relaxed.cif").write_bytes(exported.content)
        # Keep the explicit bond indices and offsets too. This is an API
        # snapshot for inspection; it is not an importable session file.
        Path("relaxed-api.json").write_text(json.dumps(relaxed, indent=2))
        print("Saved relaxed.cif and relaxed-api.json")


if __name__ == "__main__":
    main()