#!/usr/bin/env python3
"""Submit once and retrieve a durable completion. Requires: pip install httpx.

Keep --idempotency-key unchanged when restarting this script for the same job.
A local deadline stops polling; it does not cancel the server's job.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
import uuid

import httpx


async def _complete_job(
    *,
    base_url,
    api_key,
    model,
    prompt,
    idempotency_key,
    max_tokens=128,
    deadline_seconds=900,
    poll_interval=1,
):
    deadline = time.monotonic() + deadline_seconds
    body = {"model": model, "prompt": prompt, "max_tokens": max_tokens}

    def remaining():
        seconds = deadline - time.monotonic()
        if seconds <= 0:
            raise TimeoutError(
                f"Polling deadline reached; resume with Idempotency-Key {idempotency_key!r}. The server job was not cancelled."
            )
        return seconds

    async def pause(response=None):
        try:
            delay = (
                float(response.headers.get("Retry-After", poll_interval))
                if response is not None
                else poll_interval
            )
        except ValueError:
            delay = poll_interval
        await asyncio.sleep(min(max(0.01, delay), 30, remaining()))

    # Explicit close, no redirect forwarding of Authorization, no hidden retries.
    async with httpx.AsyncClient(
        base_url=base_url.rstrip("/") + "/",
        follow_redirects=False,
        headers={"Authorization": f"Bearer {api_key}"},
    ) as client:

        async def request(method, path, **kwargs):
            for attempt in range(5):
                timeout = httpx.Timeout(min(30, remaining()), connect=min(10, remaining()))
                try:
                    async with asyncio.timeout(remaining()):
                        return await client.request(method, path, timeout=timeout, **kwargs)
                except httpx.TransportError:
                    if attempt == 4:
                        raise
                    await pause()
            raise AssertionError("unreachable")

        # Retries reuse the exact body and caller-supplied key, including after a
        # lost submit acknowledgement. A different body under this key is 409.
        for attempt in range(5):
            submitted = await request(
                "POST",
                "completion-jobs",
                json=body,
                headers={"Idempotency-Key": idempotency_key},
            )
            if submitted.status_code not in (429, 500, 502, 503, 504) or attempt == 4:
                break
            await pause(submitted)
        submitted.raise_for_status()
        job_id = str(uuid.UUID(submitted.json()["id"]))
        path = f"completion-jobs/{job_id}"
        while True:
            state = await request("GET", path)
            state.raise_for_status()
            metadata = state.json()
            if metadata["status"] in ("queued", "running"):
                await pause(state)
                continue
            result = await request("GET", path + "/result")
            if result.status_code == 202:
                await pause(result)
                continue
            # Metadata HTTP 200 only means it was read. The result endpoint
            # carries success or the actual terminal failure HTTP status.
            result.raise_for_status()
            payload = result.json()
            if "error" in payload or "choices" not in payload:
                raise ValueError("Result is not a successful completion")
            return payload


def complete_job(**kwargs):
    return asyncio.run(_complete_job(**kwargs))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--base-url",
        default=os.environ.get("ACS_API_BASE", "https://infra.acsresearch.org/v1"),
    )
    parser.add_argument("--model", required=True)
    parser.add_argument("--prompt", default="Hello")
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--idempotency-key", required=True)
    parser.add_argument("--deadline", type=float, default=900)
    args = parser.parse_args()
    try:
        result = complete_job(
            base_url=args.base_url,
            api_key=os.environ["ACS_API_KEY"],
            model=args.model,
            prompt=args.prompt,
            max_tokens=args.max_tokens,
            idempotency_key=args.idempotency_key,
            deadline_seconds=args.deadline,
        )
    except (httpx.HTTPError, TimeoutError, ValueError, KeyError) as exc:
        print(f"Completion failed: {exc}", file=sys.stderr)
        return 1
    print(json.dumps(result, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())