#!/usr/bin/env python3
"""Call the faceswap-lissie RunPod serverless endpoint ASYNCHRONOUSLY (/run).

Standalone: python 3.8+, stdlib only. Same arguments as
faceswap_lissie_sync.py plus:
    --no-wait        submit only; print the job id and exit
    --job-id ID      skip submission; poll/collect an existing job
    --poll SECONDS   poll interval (default 3)
    --cancel JOB_ID  cancel a queued/running job and exit

    # fire-and-forget, collect later
    python faceswap_lissie_async.py --body-image scene.jpg --face-image face.jpg --no-wait
    python faceswap_lissie_async.py --job-id JOB_ID --out ./results
"""
import argparse
import base64
import json
import os
import pathlib
import sys
import time
import urllib.error
import urllib.request

DEFAULT_ENDPOINT = "qkeibkeoefcrk3"
TERMINAL = ("COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT")


def api(args, method, path, payload=None, timeout=90):
    req = urllib.request.Request(
        f"https://api.runpod.ai/v2/{args.endpoint_id}/{path}", method=method,
        data=json.dumps(payload).encode() if payload is not None else None,
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {args.api_key}"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.load(r)
    except urllib.error.HTTPError as e:
        sys.exit(f"ERROR: HTTP {e.code} on /{path}: {e.read().decode(errors='replace')[:2000]}")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--api-key", default=os.environ.get("RUNPOD_API_KEY"))
    ap.add_argument("--endpoint-id", default=DEFAULT_ENDPOINT)
    ap.add_argument("--body-image")
    ap.add_argument("--face-image")
    ap.add_argument("--instruction")
    ap.add_argument("--expression-prompt")
    ap.add_argument("--negative-prompt")
    ap.add_argument("--steps", type=int)
    ap.add_argument("--cfg", type=float)
    ap.add_argument("--seed", type=int)
    ap.add_argument("--rmbg-sensitivity", type=float)
    ap.add_argument("--rmbg-process-res", type=int)
    ap.add_argument("--workflow")
    ap.add_argument("--set", action="append", metavar="NODE.INPUT=VALUE")
    ap.add_argument("--workflow-json")
    ap.add_argument("--out", default=".")
    ap.add_argument("--timeout", type=float, default=1800)
    ap.add_argument("--poll", type=float, default=3.0)
    ap.add_argument("--no-wait", action="store_true")
    ap.add_argument("--job-id")
    ap.add_argument("--cancel", metavar="JOB_ID")
    args = ap.parse_args()
    if not args.api_key:
        sys.exit("ERROR: pass --api-key or set RUNPOD_API_KEY")

    if args.cancel:
        print(json.dumps(api(args, "POST", f"cancel/{args.cancel}"), indent=2))
        return

    if args.job_id:
        job_id = args.job_id
    else:
        if not args.body_image or not args.face_image:
            sys.exit("ERROR: --body-image and --face-image are required to submit a job")
        params = {"body_image": os.path.basename(args.body_image),
                  "face_image": os.path.basename(args.face_image)}
        for cli, p in [("instruction", "instruction"), ("expression_prompt", "expression_prompt"),
                       ("negative_prompt", "negative_prompt"), ("steps", "steps"),
                       ("cfg", "cfg"), ("seed", "seed"),
                       ("rmbg_sensitivity", "rmbg_sensitivity"),
                       ("rmbg_process_res", "rmbg_process_res")]:
            v = getattr(args, cli)
            if v is not None:
                params[p] = v
        payload = {
            "images": [
                {"name": os.path.basename(args.body_image),
                 "image": base64.b64encode(open(args.body_image, "rb").read()).decode()},
                {"name": os.path.basename(args.face_image),
                 "image": base64.b64encode(open(args.face_image, "rb").read()).decode()},
            ],
            "params": params,
        }
        if args.workflow:
            payload["workflow"] = args.workflow
        if args.workflow_json:
            payload["workflow_json"] = json.load(open(args.workflow_json))
        overrides = {}
        for s in args.set or []:
            k, v = s.split("=", 1)
            try:
                overrides[k] = json.loads(v)
            except json.JSONDecodeError:
                overrides[k] = v
        if overrides:
            payload["set"] = overrides
        job = api(args, "POST", "run", {"input": payload})
        job_id = job["id"]
        print(f"# submitted job {job_id}", file=sys.stderr)
        if args.no_wait:
            print(job_id)
            return

    t0 = time.monotonic()
    while True:
        result = api(args, "GET", f"status/{job_id}")
        status = result.get("status")
        if status in TERMINAL:
            break
        if time.monotonic() - t0 > args.timeout:
            sys.exit(f"ERROR: timed out after {args.timeout}s (job {job_id} still collectable with --job-id)")
        print(f"# {status} ... {time.monotonic()-t0:.0f}s", file=sys.stderr)
        time.sleep(args.poll)

    if status != "COMPLETED":
        sys.exit(f"ERROR: {json.dumps(result, indent=2)[:3000]}")
    output = result["output"]
    if "error" in output:
        sys.exit(f"ERROR from handler: {output['error']}")
    out_dir = pathlib.Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    for i, img in enumerate(output["images"]):
        p = out_dir / f"{int(time.time())}-{i}-{img['filename']}"
        p.write_bytes(base64.b64decode(img["data"]))
        print(p)
    print(f"# seed={output.get('seed')} delay={result.get('delayTime')}ms "
          f"exec={result.get('executionTime')}ms", file=sys.stderr)


if __name__ == "__main__":
    main()