# -*- coding: utf-8 -*-
"""白纸彩色板报版式（A4 横版 2896x2048）—— 不是深色黑板，是印在 A4 纸上的彩色设计稿"""
import json, os, time, requests
from concurrent.futures import ThreadPoolExecutor
from gen_chalk import LAYOUTS

CFG = json.load(open('/data/www/files/LhyjdesignPro/config.json'))
API_KEY, BASE, MODEL = CFG['api_key'], CFG['base_url'], CFG['models']['image_gen']
OUT = '/data/www/files/heibanbaoArt/rich'
os.makedirs(OUT, exist_ok=True)

STYLE = (
 "A hand-drawn COLOURFUL CLASSROOM POSTER LAYOUT (Chinese primary school 黑板报 版式设计稿) printed on a sheet "
 "of white A4 paper, seen perfectly straight-on and filling the whole frame. "
 "PAPER: clean BRIGHT WHITE paper background with a very faint warm off-white tint and subtle paper grain; a "
 "touch of empty white margin at the left, right and bottom edges. "
 "MEDIUM: a teacher's hand-painted layout template — BOLD hand-drawn outlines in dark brown-black ink, filled "
 "with FLAT BRIGHT MARKER COLOURS. Cheerful saturated school-poster palette: bright red, warm golden yellow, "
 "sky blue, grass green, orange, pink, purple, warm brown, cream. Colourful and lively, plenty of colour all "
 "over the page. "
 "BRIGHT FLAT LOOK: flat clean colour fills, NO gradients inside objects, NO shading, NO 3D, NO photorealism, "
 "NO muddy or dark colours. "
 "LINE RULE: every outline is ONE single colour — dark brown-black ink. Do NOT draw coloured outlines, do "
 "NOT draw red-and-blue offset double lines, do NOT produce any misprint / colour-registration effect. "
 "COLOUR RICHNESS: use at least SIX different bright colours (red, yellow, sky blue, grass green, orange, "
 "pink) spread across the whole page so it looks rich, cheerful and full of colour. "
 "CRITICAL — this is NOT a chalkboard: absolutely NO dark green or black background, NO chalk texture, NO "
 "white-chalk-on-black writing, NO wooden frame, NO easel, NO classroom scene. The page is WHITE PAPER with "
 "colourful drawings on it. "
 "COMPOSITION: a well-balanced poster-like page of hand-drawn panel frames plus illustrations, like a printed "
 "worksheet that a child fills in. Every panel frame is drawn with a bold coloured outline and its INTERIOR IS "
 "COMPLETELY EMPTY white paper, ready for a pupil to write into — never fill a panel interior with anything. "
 "TITLE BAND: the TOP 25 PERCENT of the page is left COMPLETELY EMPTY white paper — absolutely no frame line, "
 "no panel, no sun, no cloud, no fruit, no character, no decoration, no symbol inside that band, because a big "
 "hand-written title will be placed there. Every panel frame and every illustration must stay BELOW the 27 "
 "percent height line of the page. "
 "ABSOLUTELY NO TEXT ANYWHERE: no Chinese characters, no letters, no words, no numbers, no digits, no "
 "calligraphy, no captions, no title, no logo, no watermark, no signature, no seal. Any signboard, sign face, "
 "screen, banner, poster or book in the scene must be completely BLANK."
)


def clean(s):
    for a, b in (('chalk-drawn', 'hand-drawn'), ('chalk panel frames', 'panel frames'),
                 ('chalk panel', 'panel'), (' in chalk', ''), ('chalk ', ''), ('chalk-', '-')):
        s = s.replace(a, b)
    return s


def gen(key, layout):
    name = key + '.jpg'
    prompt = STYLE + clean(layout)
    for a in range(4):
        try:
            r = requests.post(f"{BASE}/images/generations",
                headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
                json={"model": MODEL, "prompt": prompt, "size": "2896x2048", "n": 1,
                      "watermark": False}, timeout=300)
            d = r.json()
            if d.get('data'):
                img = requests.get(d['data'][0]['url'], timeout=180).content
                p = os.path.join(OUT, name)
                open(p, 'wb').write(img)
                print(f"OK  {name}  {len(img)//1024}KB", flush=True)
                return p
            print(f"ERR {name} #{a}: {json.dumps(d, ensure_ascii=False)[:260]}", flush=True)
        except Exception as e:
            print(f"EXC {name} #{a}: {e}", flush=True)
        time.sleep(8)
    return None


if __name__ == '__main__':
    with ThreadPoolExecutor(max_workers=5) as ex:
        list(ex.map(lambda kv: gen(*kv), LAYOUTS.items()))