#!/usr/bin/env python3
"""Submit and poll asynchronous Sub2API image tasks."""
from __future__ import annotations
import argparse, base64, json, os, sys, time
from pathlib import Path
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

ORIGIN = "https://sub2apius.shujk.top"
AUTH = Path.home() / ".codex/auth.json"

def die(msg: str) -> None:
    print(f"Error: {msg}", file=sys.stderr); raise SystemExit(1)

def key(path: Path) -> str:
    try: data = json.loads(path.read_text())
    except Exception as exc: die(f"cannot read Codex auth: {exc}")
    value = data.get("OPENAI_API_KEY") if isinstance(data, dict) else None
    if not isinstance(value, str) or not value.strip(): die("OPENAI_API_KEY is missing")
    return value.strip()

def request(method: str, url: str, token: str, body: bytes | None = None) -> tuple[int, dict, dict]:
    headers = {"Authorization": f"Bearer {token}", "Accept": "application/json", "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/131 Safari/537.36"}
    if body is not None: headers["Content-Type"] = "application/json"
    req = Request(url, data=body, headers=headers, method=method)
    try:
        with urlopen(req, timeout=60) as resp:
            return resp.status, dict(resp.headers), json.loads(resp.read())
    except HTTPError as exc:
        raw = exc.read().decode(errors="replace")
        try: payload = json.loads(raw)
        except json.JSONDecodeError: payload = {"error": raw}
        return exc.code, dict(exc.headers), payload
    except URLError as exc: die(f"request failed: {exc}")

def main() -> int:
    p = argparse.ArgumentParser(description="Asynchronous Shujk image generation")
    p.add_argument("--auth-file", type=Path, default=AUTH)
    p.add_argument("--model", default="gpt-image-2")
    p.add_argument("--prompt", required=True)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--size", default="1024x1024")
    p.add_argument("--quality", default="low")
    p.add_argument("--poll-seconds", type=float, default=3)
    p.add_argument("--timeout-seconds", type=float, default=1800)
    p.add_argument("--dry-run", action="store_true")
    args = p.parse_args()
    payload = {"model": args.model, "prompt": args.prompt, "size": args.size, "quality": args.quality}
    if args.dry_run:
        print(json.dumps({"url": ORIGIN + "/v1/images/generations/async", "payload": payload}, indent=2)); return 0
    token = key(args.auth_file)
    code, headers, task = request("POST", ORIGIN + "/v1/images/generations/async", token, json.dumps(payload).encode())
    if code != 202: die(f"async submit failed ({code}): {task}")
    task_url = task.get("poll_url") or f"{ORIGIN}/v1/images/tasks/{task.get('task_id')}"
    if task_url.startswith("/"): task_url = ORIGIN + task_url
    deadline = time.monotonic() + args.timeout_seconds
    while time.monotonic() < deadline:
        try:
            code, headers, state = request("GET", task_url, token)
        except TimeoutError:
            time.sleep(args.poll_seconds)
            continue
        if code >= 400: die(f"async poll failed ({code}): {state}")
        status = state.get("status")
        if status == "completed":
            url = state.get("image_url")
            if not url: die("completed task has no image_url")
            data = None
            last_error = None
            for attempt in range(6):
                try:
                    with urlopen(Request(url, headers={"Accept": "image/*", "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/131 Safari/537.36"}), timeout=120) as resp:
                        data = resp.read()
                    break
                except HTTPError as exc:
                    last_error = exc
                    if exc.code != 403 or attempt == 5: break
                    time.sleep(2 + attempt)
                except Exception as exc:
                    last_error = exc
                    break
            if data is None: die(f"image download failed: {last_error}")
            args.out.parent.mkdir(parents=True, exist_ok=True); args.out.write_bytes(data)
            print(json.dumps({"task_id": task.get("task_id"), "status": status, "out": str(args.out), "bytes": len(data)})); return 0
        if status == "failed": die(f"image task failed: {state.get('error')}")
        delay = float(headers.get("Retry-After", args.poll_seconds)); time.sleep(max(1.0, delay))
    die("timed out waiting for image task")

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