#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""智谱 CogView 生成 5 张「案例场景」概念插画（生成阶段脚本，不写入成品页）。
色板：宝锐蓝 #1F4E9C 主 + IKB #002FA7 点缀；无文字、无真人肖像。
原始 PNG 存 _img_price_deck/ai_<name>.png
"""
import io
import json
import os
import re
import sys
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'
os.makedirs(OUT, exist_ok=True)


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) or re.match(r'^\s*provider\s*:\s*zhipu\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 and m.group(1) not in ('', "''", '""'):
                    return m.group(1)
    raise SystemExit('未找到 zhipu api_key')


KEY = zhipu_key()

PAL = ('Strict limited color palette: deep corporate blue #1F4E9C as the dominant color, '
       'electric ultramarine #002FA7 used only as small accents, off-white #FAFAF8 '
       'background, thin crisp linework')
NOTEXT = ('no text, no words, no letters, no numbers, no typography, no watermark, '
          'no signature, no logo, no captions anywhere in the image')
NOPEOPLE = ('abstract faceless geometric shapes only, no faces, no real people, '
            'no portrait, no photorealism')

PROMPTS = [
    ('scene_tug',
     'Flat vector business concept illustration, Swiss international minimalist poster. '
     'A negotiation scene: a long horizontal table in the middle; on the left and right, '
     'two opposing sides drawn as abstract faceless geometric blocks and simple rectangles '
     'only, pulling a taut rope in a tug of war across the table; bold tension lines, '
     'converging pressure lines toward the rope. ' + NOPEOPLE + '. ' + PAL + ', flat 2D, '
     'generous negative space, editorial poster composition. ' + NOTEXT + '.'),
    ('scene_weigh',
     'Flat vector business concept illustration, Swiss international minimalist poster. '
     'One single anonymous faceless geometric figure made of plain shapes standing at the '
     'center, hesitating between two options: on the left a pathway with one stack of '
     'geometric blocks, on the right a different pathway with a different stack; above '
     'them a simple balance-scale tipping slightly to one side; a dashed branching arrow '
     'splitting left and right. ' + NOPEOPLE + '. ' + PAL + ', flat 2D, thin clean lines, '
     'generous negative space. ' + NOTEXT + '.'),
    ('scene_deal',
     'Flat vector business concept illustration, Swiss international minimalist poster. '
     'Two hands clasped in a handshake at the center of the frame, drawn as extremely '
     'simplified flat geometric silhouettes with simple fingers and a clear thumb, '
     'symmetrical and clean; ' + NOPEOPLE + ' except for the two abstract hand shapes; '
     'clean converging lines pointing to the handshake and a small upward arrow. '
     + PAL + ', flat 2D vector, thin clean lines, generous negative space. ' + NOTEXT + '.'),
    ('scene_price',
     'Flat vector business concept illustration, Swiss international minimalist poster. '
     'An abstract expensive-price concept: one oversized completely blank paper price tag '
     'shape with a diagonal rejection slash, beside neat stacks of geometric coin discs; '
     'a simple seesaw beam tipping between a coin and a product cube, plus a bold downward '
     'arrow. The tag is empty and completely blank. ' + NOPEOPLE + '. ' + PAL + ', flat 2D '
     'vector, thin clean lines, generous negative space. ' + NOTEXT + '.'),
    ('jobs_cart',
     'Flat vector business concept illustration, Swiss international minimalist poster. '
     'A conceptual metaphor of progress and innovation: on the left a simple flat geometric '
     'wheeled carriage silhouette made of circles and rectangles with no animal and no '
     'horse, in the middle a bold horizontal arrow, on the right a sleek simple flat '
     'geometric car silhouette; expressing the shift from a faster horse to the motor car. '
     + NOPEOPLE + '. ' + PAL + ', flat 2D vector, thin clean lines, generous negative '
     'space. ' + NOTEXT + '.'),
]

SIZES = ['1536x1024', '1344x768']
MODELS = ['cogview-4-250304', 'cogview-4', 'cogview-3-flash']


def gen(name, prompt):
    png = os.path.join(OUT, 'ai_' + name + '.png')
    if os.path.exists(png) and os.path.getsize(png) > 20000:
        print('  %-12s 已存在，跳过 (%dKB)' % (name, os.path.getsize(png) // 1024))
        return True
    for model in MODELS:
        for size in SIZES:
            for attempt in (1, 2):
                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/%s 第%d次 空URL %s' % (model, size, attempt,
                                                            json.dumps(d, ensure_ascii=False)[:200]))
                        time.sleep(3)
                        continue
                    urllib.request.urlretrieve(url, png)
                    print('  %-12s ✅ %s @ %s  %dKB' % (name, model, size,
                                                        os.path.getsize(png) // 1024))
                    return True
                except urllib.error.HTTPError as e:
                    print('    %s/%s HTTP %s %s' % (model, size, e.code,
                                                    e.read().decode('utf-8', 'replace')[:180]))
                except Exception as e:
                    print('    %s/%s ERR %s' % (model, size, str(e)[:160]))
                time.sleep(2)
    print('  %-12s ❌ 全部失败' % name)
    return False


if __name__ == '__main__':
    print('智谱 CogView → ', OUT)
    ok = 0
    for n, p in PROMPTS:
        if gen(n, p):
            ok += 1
    print('完成 %d/%d' % (ok, len(PROMPTS)))