# -*- coding: utf-8 -*-
"""Retry remaining 4 missing outputs, skip already-downloaded, tolerate 404."""
import json, ssl, pathlib, urllib.request, hashlib, re

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()

def get(url):
    req = urllib.request.Request(url, headers={"Authorization": "Bearer " + pat})
    return urllib.request.urlopen(req, context=ctx, timeout=180)

local = {}
for p in pathlib.Path(r"D:\腾讯\吕布").rglob("*"):
    if p.is_file() and p.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp"):
        local[hashlib.sha256(p.read_bytes()).hexdigest()] = str(p)

INCLUDE = ["吕布", "赤兔", "方天画戟", "画戟", "虎牢", "濮阳", "辕门", "下邳", "凤仪亭", "董卓", "貂蝉",
           "连环计", "三英", "关羽", "张飞", "刘备", "一骑当千", "演武", "古战场", "三国", "束发",
           "凤翅", "兽面", "三视图", "翎羽", "空场景", "角弓"]
EXCLUDE = ["沈惊鸿", "狙击", "装甲", "末世", "废墟", "孙尚香"]

tasks, page = [], 1
while page <= 20:
    o = json.loads(get(f"{BASE}/api/external/v1/tasks?page={page}&page_size=100&project_id=30").read().decode("utf-8"))
    data = o.get("data") if isinstance(o.get("data"), dict) else o
    items = data.get("items") or o.get("items") or []
    if not items:
        break
    tasks.extend(items)
    total = data.get("total") or o.get("total") or 0
    if len(tasks) >= total:
        break
    page += 1

OUT = pathlib.Path(r"D:\腾讯\吕布\视频资产\后台补档")
OUT.mkdir(parents=True, exist_ok=True)

ok = fail = 0
for it in tasks:
    if ((it.get("created_at") or "")[:10]) < "2026-10-23":
        continue
    if it.get("task_type") != "image" or it.get("status") != "completed" or not it.get("image_url"):
        continue
    pr = it.get("prompt") or ""
    if any(x in pr for x in EXCLUDE) or not any(x in pr for x in INCLUDE):
        continue
    tid = it["id"]
    d = json.loads(get(f"{BASE}/api/external/v1/tasks/{tid}").read().decode("utf-8"))
    pr = (d.get("prompt") or "").replace("\n", " ")
    date = (d.get("created_at") or "")[:10].replace("-", "")
    for oi, out in enumerate(x for x in (d.get("outputs") or []) if isinstance(x, dict)):
        sha = out.get("sha256")
        url = out.get("url_snapshot") or d.get("image_url")
        if not url or (sha and (sha in local or sha in {hashlib.sha256(p.read_bytes()).hexdigest() for p in OUT.glob("*.png")})):
            continue
        snip = re.sub(r'[\\/:*?"<>|\s]+', "", pr)[:30]
        fname = f"task{tid}_o{oi}_{date}_{snip}.png"
        if (OUT / fname).exists():
            continue
        full = url if url.startswith("http") else BASE + url
        try:
            data = get(full).read()
            (OUT / fname).write_bytes(data)
            print("downloaded", fname, len(data))
            ok += 1
        except Exception as e:
            print("FAIL task", tid, "output", oi, "->", str(e)[:80])
            fail += 1
print("DONE ok =", ok, "fail =", fail)