# -*- coding: utf-8 -*-
"""用即梦 CLI 生成第 1 章 PPT 配图：逐张提交→轮询→落盘，并记录积分变化。"""
import json, os, re, subprocess, sys, time

D = '/home/zyw/.local/bin/dreamina'
BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(BASE, 'assets', 'ch1')
LOG = os.path.join(BASE, '_work', 'gen_images.log')
ENV = {k: v for k, v in os.environ.items()
       if k.lower() not in ('all_proxy', 'http_proxy', 'https_proxy', 'no_proxy')}

JOBS = [
 ('cover',      '教育课件封面插画：一台打开的笔记本电脑屏幕上显示代码与红色缺陷标记，旁边一只放大镜正在逐行检查代码，背景是深蓝色科技感渐变与细网格，扁平矢量风格，干净简洁，构图留白，无文字'),
 ('case-compat','教学插画：1990 年代风格的米白色台式电脑与一张游戏光盘盒，屏幕显示安装失败的错误提示窗口，旁边一颗发光的 CPU 芯片与错误的运算数字，扁平矢量插画，冷蓝灰色调，无文字无品牌标识'),
 ('case-mars',  '教学插画：火星探测器在下降过程中姿态失控，推进器喷出火焰，红色星球表面与橙色沙尘，背景星空，电影感写实插画风格，科技氛围，无文字'),
 ('case-medical','教学插画：医院放射治疗室内的大型放疗设备，控制屏幕显示异常的剂量曲线与红色警示，冷色调医疗环境，庄重克制，写实插画风格，画面中没有人物、没有文字'),
 ('case-xueshi','教学插画：1990 年代中国单机游戏的光盘盒与宣传海报堆叠在木桌上，纸张泛黄、印刷颗粒质感，怀旧暖色调，年代感，画面中没有真实品牌文字与标识'),
]

def log(msg):
    line = time.strftime('%H:%M:%S ') + str(msg)
    print(line, flush=True)
    open(LOG, 'a', encoding='utf-8').write(line + '\n')

def run(args, timeout=300):
    p = subprocess.run([D] + args, capture_output=True, text=True, env=ENV, timeout=timeout)
    return p.returncode, (p.stdout or '') + (p.stderr or '')

def credit():
    rc, out = run(['user_credit'], 120)
    m = re.search(r'"total_credit"\s*:\s*(\d+)', out)
    return int(m.group(1)) if m else None

c0 = credit()
log(f'起始积分 {c0}')
report = []
for name, prompt in JOBS:
    dest = os.path.join(OUT, name + '.png')
    if os.path.exists(dest) and os.path.getsize(dest) > 50000:
        log(f'[skip] {name} 已存在'); continue
    rc, out = run(['text2image', f'--prompt={prompt}', '--resolution_type=2k',
                   '--ratio=16:9', '--model_version=5.0', '--generate_num=1'], 300)
    m = re.search(r'"submit_id"\s*:\s*"([^"]+)"', out)
    if not m:
        log(f'[fail] {name} 提交失败 rc={rc} out={out[:400]}'); report.append((name, None, None)); continue
    sid = m.group(1)
    log(f'[submit] {name} submit_id={sid}')
    files = []
    seen_before = set(os.listdir(OUT))
    for i in range(40):
        time.sleep(12)
        rc, out = run(['query_result', f'--submit_id={sid}', f'--download_dir={OUT}'], 180)
        if os.path.exists(dest) and os.path.getsize(dest) > 50000:
            files = [dest]; break
        got = [f for f in os.listdir(OUT) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
        new = [f for f in got if f not in seen_before]
        if new:
            files = [os.path.join(OUT, f) for f in new]; break
    if files:
        f0 = files[0]
        if f0 != dest:
            os.replace(f0, dest)
        log(f'[done] {name} -> {dest} ({os.path.getsize(dest)} bytes)')
    else:
        log(f'[timeout] {name} 未取到结果，稍后用 list_task 找回 (submit_id={sid})')
    report.append((name, sid, dest if os.path.exists(dest) else None))
    seen_before = set(os.listdir(OUT))
c1 = credit()
log(f'结束积分 {c1} ｜ 本次消耗 {(c0 - c1) if (c0 and c1) else None}')
json.dump({'credit_before': c0, 'credit_after': c1, 'items': report},
          open(os.path.join(BASE, '_work', 'gen_images_report.json'), 'w', encoding='utf-8'),
          ensure_ascii=False, indent=2)
log('全部完成')