# -*- coding: utf-8 -*-
import json, ssl, pathlib, urllib.request, urllib.error

BASE = "https://192.168.31.88:8042"
CFG = pathlib.Path(r"C:\Users\admin\.zcode\cli\config.json")
pat = json.loads(CFG.read_text(encoding="utf-8-sig"))["mcp"]["servers"]["sorb"]["env"]["SORB_PAT"]
ctx = ssl._create_unverified_context()
OUT = pathlib.Path(r"D:\腾讯\吕布\视频资产\分镜")
OUT.mkdir(parents=True, exist_ok=True)


def get_task(tid):
    req = urllib.request.Request(BASE + "/api/external/v1/tasks/" + str(tid),
                                 headers={"Authorization": "Bearer " + pat})
    with urllib.request.urlopen(req, context=ctx, timeout=60) as r:
        return json.loads(r.read().decode("utf-8"))


def find_url(o):
    # walk common shapes
    for key in ("outputs", "output", "result", "files", "assets"):
        if key in o:
            v = o[key]
            if isinstance(v, list):
                for item in v:
                    if isinstance(item, dict):
                        for uk in ("url", "url_snapshot", "file_url", "image_url", "download_url"):
                            if item.get(uk):
                                return item[uk]
                    elif isinstance(item, str):
                        return item
            elif isinstance(v, dict):
                for uk in ("url", "url_snapshot", "file_url", "image_url", "download_url"):
                    if v.get(uk):
                        return v[uk]
    for uk in ("url", "url_snapshot", "file_url", "image_url", "download_url"):
        if o.get(uk):
            return o[uk]
    return None


for tid, name in ((13070, "S1-3v3_gpt"), (13071, "S1-3v3_seedream")):
    try:
        o = get_task(tid)
    except urllib.error.HTTPError as e:
        print(tid, "HTTP", e.code, e.read().decode("utf-8")[:500])
        continue
    print("TASK", tid, "status=", o.get("status"), "progress=", o.get("progress"))
    print(" keys=", list(o.keys()))
    url = find_url(o)
    print(" url=", url)
    if url and o.get("status") == "completed":
        full = url if url.startswith("http") else BASE + url
        req = urllib.request.Request(full, headers={"Authorization": "Bearer " + pat})
        try:
            with urllib.request.urlopen(req, context=ctx, timeout=120) as r:
                data = r.read()
            path = OUT / (name + ".png")
            path.write_bytes(data)
            print(" saved", path, len(data))
        except Exception as e:
            print(" download failed", repr(e))