#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""重生跑偏的 3 张（每张多候选，供挑选）。输出 _img_price_deck/cand_<name>_<n>.png"""
import io
import json
import os
import re
import time
import urllib.request

ENV = os.path.expanduser('~/.hermes/.env')
BASE = 'https://api.minimaxi.com/v1/image_generation'
OUT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类/培训与上手/培训/_img_price_deck'
KEY = re.search(r'^MINIMAX_CN_API_KEY=(.*)$',
                io.open(ENV, encoding='utf-8').read(), re.M).group(1).strip()

BG = 'on a plain warm off-white background with wide empty margins'
STY = 'Only flat blue shapes with thin dark outlines, flat 2D, no shading'

CANDS = [
    ('scene_tug', 1,
     'Flat vector illustration. A thick taut rope stretched horizontally across the '
     'middle of the frame. On the left end of the rope a plain flat blue square is '
     'pulling it leftwards; on the right end a plain flat blue square is pulling it '
     'rightwards; short thin tension lines beside each square. ' + STY + ' ' + BG +
     '. Nothing else in the picture.'),
    ('scene_tug', 2,
     'Flat vector illustration of a tug of war. A thick taut rope stretched across the '
     'middle of the frame. At the left end two plain flat blue rectangular blocks lean '
     'backwards pulling the rope to the left; at the right end two plain flat blue '
     'rectangular blocks lean backwards pulling the rope to the right; a small knot at '
     'the centre of the rope. ' + STY + ' ' + BG + '. Nothing else in the picture.'),
    ('jobs_cart', 1,
     'Flat vector illustration with exactly three flat blue objects arranged in a '
     'horizontal row, evenly spaced across the middle of the frame: on the left a simple '
     'old-fashioned carriage with two large round wheels, in the middle a bold thick '
     'horizontal arrow pointing to the right, on the right a simple modern motor car with '
     'two round wheels. ' + STY + ' ' + BG + '. Nothing else in the picture.'),
    ('jobs_cart', 2,
     'Flat vector illustration. On the left side of the frame a simple old-fashioned '
     'carriage with two large round wheels; on the right side of the frame a simple '
     'modern motor car with two round wheels; between them a bold thick blue arrow '
     'pointing from the carriage to the car. ' + STY + ' ' + BG + '. Nothing else in the '
     'picture.'),
    ('scene_price', 1,
     'Flat vector illustration. A large plain blank paper price tag hanging from a short '
     'string, tilted slightly, placed in the centre of the frame; beside it on the left a '
     'neat stack of three flat blue coins. The tag surface is completely blank and empty. '
     + STY + ' ' + BG + '. Nothing else in the picture.'),
]


def gen(name, n, prompt):
    png = os.path.join(OUT, 'cand_%s_%d.png' % (name, n))
    for attempt in (1, 2, 3):
        try:
            body = json.dumps({'model': 'image-01', 'prompt': prompt,
                               'image_size': '1536x1024', 'num_images': 1}).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())
            if d.get('base_resp', {}).get('status_code') not in (0, None):
                print('    %s_%d 接口错误 %s' % (name, n, d.get('base_resp')))
            url = (d.get('data') or {}).get('image_urls', [''])[0] \
                if isinstance(d.get('data'), dict) else ''
            if not url:
                time.sleep(3)
                continue
            urllib.request.urlretrieve(url, png)
            print('  %-12s #%d ✅ %dKB' % (name, n, os.path.getsize(png) // 1024))
            return True
        except Exception as e:
            print('    %s_%d ERR %s' % (name, n, str(e)[:130]))
        time.sleep(3)
    print('  %-12s #%d ❌' % (name, n))
    return False


if __name__ == '__main__':
    print('候选重生 →', OUT)
    for a, b, c in CANDS:
        gen(a, b, c)