Adom Chip Thumbnailer

#!/usr/bin/env python3
"""Client for the shared Adom STEP service (service-step2glb): /step-meta, /laser-etch, /thumbnail and
/outline, sync (200) or async (202 + job, polled). Stdlib only. Every POST names its caller with X-Client
and X-Job-Name, as the service asks."""
import json, os, time, urllib.error, urllib.parse, urllib.request

URL = os.environ.get("STEP2GLB_SERVICE_API", "https://step2glb-gmdoncpxdwx0.adom.cloud").rstrip("/")
CLIENT = os.environ.get("AIFLOW_ETCH_CLIENT", "adom-aiflow-etch/" + os.environ.get("USER", "user"))


def _get(path, job, timeout=60):
    req = urllib.request.Request(URL + path, headers={"X-Client": CLIENT, "X-Job-Name": job})
    return urllib.request.urlopen(req, timeout=timeout)


def call(path, body, job, retries=2):
    """POST a STEP (bytes or str) to `path` and return (bytes, headers)."""
    if isinstance(body, str):
        body = body.encode("latin-1")
    last = None
    for attempt in range(retries + 1):
        try:
            req = urllib.request.Request(URL + path, data=body, method="POST", headers={
                "Content-Type": "application/step", "X-Client": CLIENT, "X-Job-Name": job[:120]})
            r = urllib.request.urlopen(req, timeout=240)
            data, hdr = r.read(), dict(r.headers)
            if r.status != 202:
                return data, hdr
            jid = json.loads(data)["job_id"]
            for _ in range(240):
                v = json.loads(_get(f"/jobs/{jid}", "poll", 30).read())
                st = v.get("status")
                if st in ("complete", "done", "succeeded"):
                    r2 = _get(f"/jobs/{jid}/result", "result", 120)
                    return r2.read(), dict(r2.headers)
                if st in ("failed", "error"):
                    raise RuntimeError(f"service job {jid} failed: {v.get('error')}")
                time.sleep(1.5)
            raise RuntimeError(f"service job {jid} timed out")
        except urllib.error.HTTPError as e:
            last = RuntimeError(f"POST {path} -> HTTP {e.code}: {e.read()[:300]!r}")
            if e.code < 500:
                raise last
        except (urllib.error.URLError, TimeoutError, ConnectionError) as e:
            last = RuntimeError(f"POST {path}: {e}")
        time.sleep(3 * (attempt + 1))
    raise last


def q(**kw):
    return urllib.parse.urlencode({k: v for k, v in kw.items() if v is not None})


def step_meta(step, job="mark-plan-meta"):
    return json.loads(call("/step-meta", step, job)[0])


def laser_etch(step, line1, line2=None, job="mark-etch"):
    return call("/laser-etch?" + q(name=line1, line2=line2, mode="flat"), step, job)[0].decode("latin-1")


def thumbnail(step, pose="top", w=640, h=640, bg=None, name=None, rot=None, job="mark-check-view"):
    data, hdr = call("/thumbnail?" + q(pose=pose, kind="step", upAxis="asIs", width=w, height=h, bg=bg, name=name, rot=rot), step, job)
    return data, hdr


def outline(step, name=None, rot=None, suffix=None, line2=None, job="mark-outline"):
    return json.loads(call("/outline?" + q(upAxis="z", name=name, line2=line2, rot=rot, suffix=suffix), step, job)[0])