#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""run-batch.py — 即梦出图批量执行器（并发 + 断点续跑 + 逐条日志）

用法:
  python3 run-batch.py prompts/01-风格测试.json [更多manifest...] [--workers 2] [--force]

manifest 为 JSON 数组，元素:
  {"outdir":"01-风格测试","name":"语义名","ratio":"16:9","prompt":"提示词"}

行为:
  - 若 <outdir>/<name>.png 已存在则跳过（--force 强制重跑）
  - 每个条目调用本目录 gen.sh，日志写 logs/<name>.log
  - 汇总每条的结果、submit_id 与积分变化
"""
import json, os, subprocess, sys, time, threading, queue

ROOT = os.path.dirname(os.path.abspath(__file__))
os.chdir(ROOT)


def credit():
    try:
        out = subprocess.run(["dreamina", "user_credit"], capture_output=True, text=True, timeout=60).stdout
        return json.loads(out[out.index("{"):out.rindex("}") + 1]).get("total_credit")
    except Exception:
        return None


def main():
    args, workers, force = [], 2, False
    argv = sys.argv[1:]
    i = 0
    while i < len(argv):
        a = argv[i]
        if a == "--force":
            force = True
        elif a == "--workers":
            i += 1
            workers = int(argv[i])
        elif a.startswith("--workers="):
            workers = int(a.split("=", 1)[1])
        else:
            args.append(a)
        i += 1

    items = []
    for mf in args:
        with open(mf, encoding="utf-8") as f:
            for it in json.load(f):
                it["_manifest"] = os.path.basename(mf)
                items.append(it)

    c0 = credit()
    print(f"[i] 待处理 {len(items)} 条，并发 {workers}，起始积分 {c0}", flush=True)

    q = queue.Queue()
    results = []
    lock = threading.Lock()

    def work():
        while True:
            try:
                it = q.get_nowait()
            except queue.Empty:
                return
            outdir, name = it["outdir"], it["name"]
            dest = os.path.join(outdir, name + ".png")
            if os.path.exists(dest) and not force:
                with lock:
                    results.append((name, "SKIP", "已存在"))
                print(f"[skip] {name}", flush=True)
                q.task_done()
                continue
            os.makedirs(outdir, exist_ok=True)
            pf = os.path.join("prompts", "_tmp_%s.txt" % name)
            with open(pf, "w", encoding="utf-8") as f:
                f.write(it["prompt"])
            logf = os.path.join("logs", name + ".log")
            t0 = time.time()
            with open(logf, "w", encoding="utf-8") as lf:
                p = subprocess.run(
                    ["bash", "gen2.sh", "--name", name, "--prompt-file", pf,
                     "--outdir", outdir, "--ratio", it.get("ratio", "16:9"),
                     "--model", it.get("model", "5.0"), "--res", it.get("res", "2k"),
                     "--num", str(it.get("num", 1))],
                    stdout=lf, stderr=subprocess.STDOUT, text=True,
                )
            os.remove(pf)
            ok = p.returncode == 0 and os.path.exists(dest)
            dt = int(time.time() - t0)
            with lock:
                results.append((name, "OK" if ok else "FAIL", f"{dt}s"))
            print(f"[{'ok' if ok else 'FAIL'}] {name} ({dt}s)", flush=True)
            q.task_done()

    for it in items:
        q.put(it)
    ths = [threading.Thread(target=work, daemon=True) for _ in range(workers)]
    [t.start() for t in ths]
    [t.join() for t in ths]

    c1 = credit()
    print("\n===== 汇总 =====")
    for n, s, d in results:
        print(f"{s:4} {n} {d}")
    used = (c0 - c1) if (c0 is not None and c1 is not None) else "?"
    print(f"积分: {c0} -> {c1} (本次消耗 {used})")
    bad = [r for r in results if r[1] == "FAIL"]
    print(f"成功 {sum(1 for r in results if r[1] == 'OK')} / 跳过 {sum(1 for r in results if r[1]=='SKIP')} / 失败 {len(bad)}")
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())