#!/usr/bin/env python3
"""縦長4コマを Gemini API のバッチで頼む（2026-11-05 エイイチ「Gemini 代を節約。バッチで多めの下書きを何話も一気に頼んでから修正」）

使い方（道具/ で動かす）:
  python3 vbatch.py submit 138 139 140 [-n 2] [-p 1 3]   … 話ごとに prompts を作り、1コマ n 枚をまとめて頼む（-p でコマを絞る）
  python3 vbatch.py status                                 … 頼んだバッチの状態
  python3 vbatch.py fetch                                  … 終わったバッチを回収し、各話の 生成/ に置く（ふつうの gen と同じ名前の形）
直しは今までどおり vbuild.py の gen（-n 1）でその場で頼む。

- 料金はふつうの頼み方の約半分。返りは最大24時間（48時間で失効）
- 参照画像は Files API に上げて fileUri で指す（同じ絵は48時間のうち使い回す。記録は _batch/files.json）
- 頼んだ中身と結果の記録は 道具/_batch/ に残す
"""
import sys, os, json, time, datetime, mimetypes, urllib.request, urllib.error, base64, glob
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import vbuild as V

API = "https://generativelanguage.googleapis.com"
BD = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_batch")
os.makedirs(BD, exist_ok=True)
FILES = BD + "/files.json"
JOBS = BD + "/jobs.json"
KEY = next(open(p).read().strip() for p in V.KEYS if os.path.exists(p))

def call(method, url, body=None, headers=None, raw=False, timeout=300):
    h = {"x-goog-api-key": KEY}
    h.update(headers or {})
    data = body if isinstance(body, (bytes, type(None))) else json.dumps(body).encode()
    if data is not None and "Content-Type" not in h: h["Content-Type"] = "application/json"
    req = urllib.request.Request(url, data=data, headers=h, method=method)
    for attempt in range(4):
        try:
            r = urllib.request.urlopen(req, timeout=timeout)
            return r if raw else json.load(r)
        except urllib.error.HTTPError as e:
            msg = e.read().decode()[:1500]
            if e.code in (429, 500, 503) and attempt < 3:
                time.sleep(10 * (attempt + 1)); continue
            raise SystemExit("失敗 %s %s\n%s" % (e.code, url.split("?")[0], msg))

def upload(path, mime, name):
    """Files API に上げる（resumable）。files/xxx の名前と uri を返す"""
    data = open(path, "rb").read()
    r = call("POST", API + "/upload/v1beta/files", {"file": {"display_name": name[:120]}}, raw=True,
             headers={"X-Goog-Upload-Protocol": "resumable", "X-Goog-Upload-Command": "start",
                      "X-Goog-Upload-Header-Content-Length": str(len(data)), "X-Goog-Upload-Header-Content-Type": mime})
    up = r.headers["X-Goog-Upload-URL"]
    info = call("POST", up, data, headers={"Content-Length": str(len(data)), "X-Goog-Upload-Offset": "0",
                                           "X-Goog-Upload-Command": "upload, finalize", "Content-Type": mime})
    return info["file"]

def ref_uri(path, cache):
    """参照画像の fileUri。中身が同じで上げてから44時間以内なら使い回す"""
    st = os.stat(path); k = os.path.abspath(path)
    c = cache.get(k)
    if c and c["mtime"] == st.st_mtime and time.time() - c["at"] < 44 * 3600:
        return c["uri"], c["mime"]
    mime = mimetypes.guess_type(path)[0] or "image/png"
    f = upload(path, mime, os.path.basename(path))
    cache[k] = {"uri": f["uri"], "name": f["name"], "mime": mime, "mtime": st.st_mtime, "at": time.time()}
    json.dump(cache, open(FILES, "w"), ensure_ascii=False, indent=0)
    return f["uri"], mime

def make_prompts(ep, m):
    d = V.epdir(ep); os.makedirs(d + "/生成", exist_ok=True)
    for i, p in enumerate(m.PANELS, 1):
        p["_guide"] = d + "/生成/_guide_P%d.png" % i
        txt, refs = V.build(p, ep)
        open(d + "/第%d話_P%d_プロンプト.txt" % (ep, i), "w", encoding="utf-8").write(txt)
        json.dump(refs, open(d + "/第%d話_P%d_refs.json" % (ep, i), "w"), ensure_ascii=False, indent=0)
    for k, w in getattr(m, "WINDOWS", {}).items():
        txt, refs = V.build_window(w)
        open(d + "/第%d話_P%d_小窓_プロンプト.txt" % (ep, k), "w", encoding="utf-8").write(txt)
        json.dump(refs, open(d + "/第%d話_P%d小窓_refs.json" % (ep, k), "w"), ensure_ascii=False, indent=0)

def submit(eps, n, which):
    cache = json.load(open(FILES)) if os.path.exists(FILES) else {}
    stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    lines, index = [], {}
    for ep in eps:
        m = V.load(ep); make_prompts(ep, m)
        for prompt_path, out_head, refs, aspect in V.jobs(ep, m, which):
            parts = []
            for label, path in refs:
                uri, mime = ref_uri(path, cache)
                parts.append({"text": "Reference image — %s:" % label})
                parts.append({"fileData": {"mimeType": mime, "fileUri": uri}})
            parts.append({"text": open(prompt_path, encoding="utf-8").read()})
            req = {"contents": [{"role": "user", "parts": parts}],
                   "generationConfig": {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": aspect, "imageSize": "1K"}}}
            for c in range(1, n + 1):
                key = "%s#%d" % (os.path.relpath(out_head, V.BASE), c)
                lines.append(json.dumps({"key": key, "request": req}, ensure_ascii=False))
                index[key] = {"prompt": prompt_path, "refs": ["%s=%s" % (a, b) for a, b in refs]}
        print("第%d話 プロンプトを作った" % ep)
    jl = BD + "/%s_input.jsonl" % stamp
    open(jl, "w", encoding="utf-8").write("\n".join(lines) + "\n")
    f = upload(jl, "application/jsonl", "manga_batch_%s" % stamp)
    res = call("POST", API + "/v1beta/models/%s:batchGenerateContent" % V.MODEL,
               {"batch": {"display_name": "manga_%s" % stamp, "input_config": {"file_name": f["name"]}}})
    name = res.get("name") or res.get("metadata", {}).get("name")
    jobs = json.load(open(JOBS)) if os.path.exists(JOBS) else []
    jobs.append({"name": name, "stamp": stamp, "eps": eps, "n": n, "count": len(lines), "input": jl, "index": index, "fetched": False})
    json.dump(jobs, open(JOBS, "w"), ensure_ascii=False, indent=1)
    print("頼んだ:", name, "枚数", len(lines), "（第%s話・1コマ%d枚）" % ("・".join(map(str, eps)), n))

def get(name):
    return call("GET", API + "/v1beta/" + name)

def state_of(s):
    md = s.get("metadata", {})
    return md.get("state") or s.get("state"), md.get("batchStats") or s.get("batchStats") or {}

def status():
    jobs = json.load(open(JOBS)) if os.path.exists(JOBS) else []
    for j in jobs:
        if j.get("fetched"): continue
        st, stats = state_of(get(j["name"]))
        print(j["stamp"], j["name"], "第%s話" % "・".join(map(str, j["eps"])), j["count"], "枚", st, json.dumps(stats, ensure_ascii=False))

def fetch():
    jobs = json.load(open(JOBS)) if os.path.exists(JOBS) else []
    for j in jobs:
        if j.get("fetched"): continue
        s = get(j["name"]); st, stats = state_of(s)
        if st not in ("JOB_STATE_SUCCEEDED", "BATCH_STATE_SUCCEEDED"):
            print(j["name"], st, "まだ回収できない"); continue
        resp = s.get("response", {})
        rf = resp.get("responsesFile") or s.get("metadata", {}).get("output", {}).get("responsesFile") or s.get("dest", {}).get("fileName")
        out = BD + "/%s_output.jsonl" % j["stamp"]
        r = call("GET", API + "/download/v1beta/%s:download?alt=media" % rf, raw=True, timeout=600)
        open(out, "wb").write(r.read())
        ok = ng = 0; tok_in = tok_out = 0
        fstamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        for line in open(out, encoding="utf-8"):
            if not line.strip(): continue
            o = json.loads(line); key = o.get("key", "")
            head, c = key.rsplit("#", 1)
            out_head = os.path.join(V.BASE, head)
            res = o.get("response") or {}
            img = None
            for cd in res.get("candidates", []):
                for p in cd.get("content", {}).get("parts", []):
                    if "inlineData" in p: img = p["inlineData"]
            base = "%s_%sb%s" % (out_head, fstamp, c)     # 例 第138話_P1_20261106_070000b1.jpg（latest・measure の _2* に合う）
            u = res.get("usageMetadata", {}); tok_in += u.get("promptTokenCount", 0); tok_out += u.get("candidatesTokenCount", 0)
            if img:
                ext = ".jpg" if "jpeg" in img["mimeType"] else ".png"
                open(base + ext, "wb").write(base64.b64decode(img["data"])); ok += 1
            else:
                ng += 1; print("画像なし", key, json.dumps(o.get("error") or res.get("candidates", [{}])[0].get("finishReason"), ensure_ascii=False)[:300])
            for cd in res.get("candidates", []):
                for p in cd.get("content", {}).get("parts", []): p.pop("inlineData", None)
            meta = dict(j["index"].get(key, {}), model=V.MODEL, batch=j["name"], response=res, error=o.get("error"))
            json.dump(meta, open(base + ".json", "w", encoding="utf-8"), ensure_ascii=False, indent=1)
        j["fetched"] = True; j["result"] = {"ok": ok, "ng": ng, "input_tokens": tok_in, "output_tokens": tok_out, "fetched_at": fstamp}
        json.dump(jobs, open(JOBS, "w"), ensure_ascii=False, indent=1)
        print(j["name"], "回収", ok, "枚・画像なし", ng, "・入力トークン", tok_in, "・出力トークン", tok_out)

if __name__ == "__main__":
    a = sys.argv[1:]
    if not a: print(__doc__); sys.exit()
    cmd, rest = a[0], a[1:]
    if cmd == "submit":
        n = 2; which = []
        if "-n" in rest: k = rest.index("-n"); n = int(rest[k + 1]); rest = rest[:k] + rest[k + 2:]
        if "-p" in rest: k = rest.index("-p"); which = rest[k + 1:]; rest = rest[:k]
        submit([int(x) for x in rest], n, which)
    elif cmd == "status": status()
    elif cmd == "fetch": fetch()
