#!/usr/bin/env python3
"""Generate images through QDivZero; Python 3.11+, no packages required.

Export QDIV0_API_KEY and QDIV0_ACCOUNT_ID before running. This sends one
billable request; it deliberately does not retry timeouts automatically.
"""

import argparse
import base64
import json
import os
from pathlib import Path
import time
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


def save_images(response, output):
    images = response.get("data", [])
    if not images:
        raise ValueError("The API returned no images.")
    saved = []
    for index, item in enumerate(images, start=1):
        data = base64.b64decode(item["b64_json"], validate=True)
        if data.startswith(b"\x89PNG\r\n\x1a\n"):
            extension = ".png"
        elif data.startswith(b"\xff\xd8\xff"):
            extension = ".jpg"
        elif data.startswith(b"RIFF") and data[8:12] == b"WEBP":
            extension = ".webp"
        else:
            raise ValueError("The response is not a supported PNG, JPEG or WebP image.")
        path = Path(f"{output}-{index}{extension}")
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("xb") as file:
            file.write(data)
        saved.append(str(path))
    return saved


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="QDivZero serving name")
    parser.add_argument("--example", choices=["product", "poster", "poster_refined", "transparent", "cover"], default="product")
    parser.add_argument("--prompts", type=Path, default=Path(__file__).with_name("prompts.json"))
    parser.add_argument("--steps", type=int, choices=range(1, 101), metavar="1..100", required=True)
    parser.add_argument("--size", default="1024x1024")
    parser.add_argument("--guidance", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", type=Path, required=True, help="New file prefix, e.g. results/teapot")
    parser.add_argument("--timeout", type=float, default=900)
    args = parser.parse_args()

    api_key = os.environ.get("QDIV0_API_KEY")
    account_id = os.environ.get("QDIV0_ACCOUNT_ID")
    if not api_key or not account_id:
        parser.error("Set QDIV0_API_KEY and QDIV0_ACCOUNT_ID first.")
    if args.seed < 0 or args.timeout <= 0:
        parser.error("Seed must be non-negative and timeout must be positive.")

    prompts = json.loads(args.prompts.read_text(encoding="utf-8"))
    payload = {
        "model": args.model,
        "prompt": prompts[args.example],
        "size": args.size,
        "n": 1,
        "num_inference_steps": args.steps,
        "guidance_scale": args.guidance,
        "seed": args.seed,
        "response_format": "b64_json",
    }
    metadata_path = Path(f"{args.output}.json")
    if metadata_path.exists() or any(args.output.parent.glob(f"{args.output.name}-*.*")):
        parser.error("Output already exists; choose a new --output prefix.")
    request = Request(
        "https://api.qdiv0.com/v1/images/generations",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "X-Account-Id": account_id,
            "Content-Type": "application/json",
        },
        method="POST",
    )
    started = time.monotonic()
    try:
        with urlopen(request, timeout=args.timeout) as response:
            result = json.load(response)
        elapsed = time.monotonic() - started
        files = save_images(result, args.output)
    except HTTPError as error:
        raise SystemExit(f"HTTP {error.code}: {error.read(2000).decode('utf-8', errors='replace')}") from None
    except (URLError, TimeoutError) as error:
        raise SystemExit(f"Request failed: {error}. Check the deployment before sending another request; the GPU may still be working.") from None
    metadata = {"request": payload, "request_seconds": round(elapsed, 3), "files": files}
    metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()