#!/usr/bin/env python3
"""T2 全量 300 张放量 — qwen21 896x1344, 双 worker 并发, 失败重试≤2, 断点续跑.
已收编的 first20 (20张) 自动 skip; 剩 280 张.
下载 = 响应 url 为准 (1018 实录: 禁写死实例端口)."""
import json, time, os, threading, queue as tq, urllib.request, csv

ROOT = "/home/kai/shared/virtual-idol/week1"
MD = f"{ROOT}/T2_matrix"
OUT = f"{MD}/full300"
os.makedirs(OUT, exist_ok=True)
os.chdir(ROOT)
LEDGER = open("manifest/t2_ledger.jsonl", "a")

rows = {int(r["combo_id"]): r for r in csv.DictReader(open(f"{MD}/matrix_300.csv"))}
all_ids = sorted(rows)

def exists(cid, r):
    # 已收编位置两处: first20/ 与 full300/
    name = f"T2_{cid:03d}_{r['A']}_{r['B']}_{r['C']}.png"
    for d in (f"{MD}/first20", OUT):
        p = os.path.join(d, name)
        if os.path.exists(p) and os.path.getsize(p) > 100_000:
            return p
    return None

jobs = [rows[i] for i in all_ids if not exists(i, rows[i])]
log_done = len(all_ids) - len(jobs)

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)
def led(rec):
    rec["ts"] = time.time()
    LEDGER.write(json.dumps(rec, ensure_ascii=False) + "\n"); LEDGER.flush()

WORK = tq.Queue()
for r in jobs: WORK.put(r)
LOCK = threading.Lock()
stats = {"ok": 0, "fail": 0, "t0": time.time()}

def gen_one(name, r, attempt=1):
    cid = int(r["combo_id"])
    out = f"{OUT}/T2_{cid:03d}_{r['A']}_{r['B']}_{r['C']}.png"
    body = json.dumps({"prompt": r["prompt_full"], "seed": cid, "width": 896, "height": 1344}).encode()
    req = urllib.request.Request("http://127.0.0.1:10588/api/production/qwen21/generate",
        data=body, headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=600) as resp:
            j = json.load(resp)
        src = j["data"]["images"][0].get("url") or ""
        src = (src or f"http://127.0.0.1:8188/view?filename={j['data']['images'][0]['filename']}&subfolder=&type=output").replace("localhost", "127.0.0.1")
        urllib.request.urlretrieve(src, out)
        el = time.time() - t0
        log(f"[{name}] {cid}({r['A']}_{r['B']}_{r['C']}): {el:.0f}s {os.path.getsize(out)//1024}KB" + (f" retry{attempt}" if attempt > 1 else ""))
        led({"event": "gen", "phase": "T2-full300", "asset_id": f"T2_{cid:03d}",
             "engine": "qwen21", "mode": "t2i", "seed": cid,
             "params": {"width": 896, "height": 1344}, "attempt": attempt,
             "final_prompt": r["prompt_full"], "output_path": os.path.abspath(out),
             "elapsed": round(el, 1), "status": "ok"})
        return True
    except Exception as e:
        log(f"[{name}] {cid}: FAIL(a{attempt}) {str(e)[:150]}")
        if attempt < 2:
            time.sleep(5)
            return gen_one(name, r, attempt + 1)
        led({"event": "gen", "phase": "T2-full300", "asset_id": f"T2_{cid:03d}",
             "engine": "qwen21", "mode": "t2i", "seed": cid,
             "final_prompt": r["prompt_full"], "status": "fail", "err": str(e)[:200]})
        return False

def worker(name):
    while True:
        try: r = WORK.get_nowait()
        except tq.Empty: return
        cid = int(r["combo_id"])
        if exists(cid, r):
            with LOCK: stats["ok"] += 1
            continue
        if gen_one(name, r):
            with LOCK: stats["ok"] += 1
        else:
            with LOCK: stats["fail"] += 1
        WORK.task_done()

threads = [threading.Thread(target=worker, args=(f"w{i}",)) for i in range(2)]
for t in threads: t.start()
for t in threads: t.join()
el_total = time.time() - stats["t0"]
log(f"FULL300_DONE ok={stats['ok']}/{len(all_ids)} fail={stats['fail']} elapsed={el_total/60:.0f}min")