#!/usr/bin/env python3
"""Exercise a deployed RunPod serverless endpoint.

    python tests/endpoint_test.py --endpoint-id XXXX main.jpg [ref.png] \
        [--prompt "..."] [--mode turbo-8] [--seed 1] [--sync]

Reads RUNPOD_API_KEY from the environment. Saves returned images next to this
script and prints per-phase timings (delay/queue vs execution).
"""
import argparse
import base64
import json
import os
import pathlib
import sys
import time
import urllib.request


def api(method, url, payload=None, timeout=120):
    req = urllib.request.Request(
        url, method=method,
        data=json.dumps(payload).encode() if payload is not None else None,
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {os.environ['RUNPOD_API_KEY']}"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("images", nargs="+")
    ap.add_argument("--endpoint-id", required=True)
    ap.add_argument("--prompt", default="replace the outfit of the subject in the first "
                    "image with the outfit in the second image, don't change anything else.")
    ap.add_argument("--mode", default=None, help="only for workflows that support it")
    ap.add_argument("--seed", type=int)
    ap.add_argument("--sync", action="store_true", help="use /runsync instead of /run+poll")
    ap.add_argument("--timeout", type=float, default=900)
    args = ap.parse_args()

    imgs = [{"name": os.path.basename(p),
             "image": base64.b64encode(open(p, "rb").read()).decode()}
            for p in args.images]
    params = {"prompt": args.prompt}
    if args.mode:
        params["mode"] = args.mode
    if args.seed is not None:
        params["seed"] = args.seed
    payload = {"input": {"images": imgs, "params": params}}
    base = f"https://api.runpod.ai/v2/{args.endpoint_id}"

    t0 = time.monotonic()
    if args.sync:
        result = api("POST", f"{base}/runsync", payload, timeout=args.timeout)
    else:
        job = api("POST", f"{base}/run", payload)
        print(f"job {job['id']} -> {job['status']}", file=sys.stderr)
        while True:
            result = api("GET", f"{base}/status/{job['id']}")
            if result["status"] in ("COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"):
                break
            if time.monotonic() - t0 > args.timeout:
                sys.exit(f"timed out after {args.timeout}s (status {result['status']})")
            print(f"  {result['status']} ... {time.monotonic()-t0:.0f}s", file=sys.stderr)
            time.sleep(2)
    wall = time.monotonic() - t0

    if result.get("status") != "COMPLETED":
        sys.exit(f"FAILED: {json.dumps(result, indent=2)[:3000]}")
    out = result["output"]
    if "error" in out:
        sys.exit(f"handler error: {out['error']}")
    for i, img in enumerate(out["images"]):
        dest = pathlib.Path(__file__).parent / f"endpoint-out-{i}-{img['filename']}"
        dest.write_bytes(base64.b64decode(img["data"]))
        print(f"saved {dest}")
    print(f"status=COMPLETED wall={wall:.1f}s delayTime={result.get('delayTime')}ms "
          f"executionTime={result.get('executionTime')}ms\n"
          f"handler timings={out.get('timings')} seed={out.get('seed')}")


if __name__ == "__main__":
    main()