#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""智谱 CogView-4 重出案例插画（v3）。
v2 教训：
  ① prompt 里写 `#1F4E9C`/`#002FA7` → 模型把色号**当文字画出来**（实测出现
     「#FI1093C ⑦FA F4DAF7」）。颜色只用词说，绝不写 #hex。
  ② 写「no text / no letters / no logo」这类否定词 → 反而诱导出字母涂鸦与角标 logo。
     改为纯正向描述 + 「背景为纯色空白」。
原始 PNG 存 _img_price_deck/ai_<name>.png
"""
import io
import json
import os
import re
import time
import urllib.request

CFG = os.path.expanduser('~/.hermes/profiles/zone-a/config.yaml')
BASE = 'https://open.bigmodel.cn/api/paas/v4/images/generations'
OUT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类/培训与上手/培训/_img_price_deck'

BG = ('A plain uniform empty off-white background, wide empty margins all around, '
      'nothing else in the picture')
STYLE = ('Minimalist flat vector illustration, only two flat colours: a deep navy blue '
         'and a medium blue, plus thin dark outlines, flat 2D, no shading, no gradient, '
         'clean and simple')

PROMPTS = [
    ('scene_tug',
     'Minimalist flat vector illustration. ' + BG + '. In the middle a long simple '
     'rectangular table. On the left side a plain dark navy rectangle leaning to the '
     'right; on the right side a plain medium blue rectangle leaning to the left; both '
     'hold a thick taut rope that runs horizontally across the table between them. '
     + STYLE + '.'),

    ('scene_weigh',
     'Minimalist flat vector illustration. ' + BG + '. A simple balance scale with a '
     'central vertical post and two hanging pans; the left pan holds one large plain blue '
     'cube, the right pan holds two small plain blue cubes; the beam tilts down to the '
     'left. ' + STYLE + '.'),

    ('scene_deal',
     'Minimalist flat vector illustration. ' + BG + '. At the centre, two very simple '
     'flat blue hands clasped together in a handshake, clean minimal shapes, five simple '
     'fingers on each hand, seen from the side. ' + STYLE + '.'),

    ('scene_price',
     'Minimalist flat vector illustration. ' + BG + '. A large plain blank paper tag '
     'shape tilted slightly with a thick diagonal line crossing it, and a small neat '
     'stack of plain blue discs beside it; below, a simple seesaw beam with a blue disc '
     'on its left end and a plain cube on its right end. The tag surface is completely '
     'plain and empty. ' + STYLE + '.'),

    ('jobs_cart',
     'Minimalist flat vector illustration. ' + BG + '. On the left a simple flat blue '
     'carriage with two round wheels, in the middle a bold horizontal arrow, on the right '
     'a simple flat blue motor car with two round wheels. ' + STYLE + '.'),
]

MODEL = 'cogview-4-250304'
SIZE = '1536x1024'


def gen(name, prompt):
    png = os.path.join(OUT, 'ai_' + name + '.png')
    for attempt in (1, 2, 3):
        try:
            body = json.dumps({'model': MODEL, 'prompt': prompt, 'size': SIZE}).encode()
            req = urllib.request.Request(BASE, data=body, method='POST',
                                         headers={'Authorization': 'Bearer ' + KEY,
                                                  'Content-Type': 'application/json'})
            with urllib.request.urlopen(req, timeout=180) as r:
                d = json.loads(r.read())
            url = (d.get('data') or [{}])[0].get('url')
            if not url:
                print('    %s 第%d次 空URL' % (name, attempt))
                time.sleep(3)
                continue
            urllib.request.urlretrieve(url, png)
            print('  %-12s ✅ %dKB' % (name, os.path.getsize(png) // 1024))
            return True
        except urllib.error.HTTPError as e:
            print('    %s HTTP %s %s' % (name, e.code, e.read().decode('utf-8', 'replace')[:160]))
        except Exception as e:
            print('    %s ERR %s' % (name, str(e)[:140]))
        time.sleep(3)
    print('  %-12s ❌' % name)
    return False


def zhipu_key():
    lines = io.open(CFG, encoding='utf-8').read().splitlines()
    for i, l in enumerate(lines):
        if re.match(r'^\s*zhipu\s*:\s*$', l):
            for j in range(i, min(i + 8, len(lines))):
                m = re.match(r'^\s*api_key\s*:\s*[\'"]?([^\'"\s]+)', lines[j])
                if m:
                    return m.group(1)
    raise SystemExit('未找到 zhipu api_key')


KEY = zhipu_key()
if __name__ == '__main__':
    print('CogView-4 v3 →', OUT)
    print('完成 %d/%d' % (sum(1 for n, p in PROMPTS if gen(n, p)), len(PROMPTS)))