# -*- coding: utf-8 -*-
"""第2步：以线条图为参考，按各主题配色方案生成配色图（图生图 strength=0.6）
严禁改动构图；仅平涂上色。每个主题单独一段配色方案，不写就是随机配色。"""
import base64, glob, json, os, sys, time, requests
from concurrent.futures import ThreadPoolExecutor

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

BASE_INSTR = ("Create the watercolour COLOURED version of this exact line drawing, painted in the style of a "
              "schoolchild's watercolour painting. Keep EXACTLY the same shapes, positions, proportions and "
              "composition — do not add, remove, move or redraw any element. Simply fill the outlined areas with "
              "clean transparent flat watercolour washes that stay neatly inside the black outlines, with the "
              "black outlines still clearly visible on top. Gentle layered washes, soft transparent pooling, "
              "visible watercolour paper grain. NO text of any kind, no letters, no characters, no numbers, no "
              "signature, no watermark, no stamp. Keep it simple and neatly painted — no heavy shading, no dark "
              "rendering, no gradients into black. ")

COLORS = {
 '1_饮水思源_一滴水的旅行':
   " Colour scheme: clear pale blue sky with two soft white clouds and a pale yellow sun with soft rays; the big "
   "mountain fresh blue-green with darker green pine trees; the long winding S-shaped water ribbon in bright "
   "clear cyan-blue with paler highlights and a few white sparkle dots; the river the same cool blue with white "
   "ripple lines; the small wooden boat warm brown with a pale cream rim stripe; the riverside reeds warm "
   "yellow-green; the stone bridge and the plaza paving warm light grey; the tall stone monument warm beige with "
   "a slightly darker stepped plaza and small pale grey lamp posts; the house wall clean white with a pale blue "
   "window; the stove dark grey metal with the kettle in warm copper and a pale cream white curly steam line; "
   "the porcelain cooling jar white with a pale blue rim; the wall tap pale grey with a small red handle; the "
   "small bucket pale blue with a white rim; the watering can light green; the flowers small red and yellow; the "
   "flower bed green; the child's shirt pale sky blue, the neckerchief pure bright red, hair dark brown; the "
   "big floating water drop bright cyan with a white highlight; the grass tufts yellow-green; the birds pale "
   "grey-blue.",

 '2_竹节里的八杯水':
   " Colour scheme: sky soft blue with two white clouds and one big warm yellow sun with pale rays; the mountain "
   "fresh green with darker pine greens and a soft white band of cloud sea across it; the thick bamboo stalk "
   "fresh yellow-green with slightly darker green rings at each of the eight nodes and pale green fibre lines; "
   "the bamboo leaf clusters in two or three tones of green; the eight small cups pale cream with pale blue "
   "water inside, the water levels getting deeper blue toward the top of the stalk; the single water drops "
   "bright cyan with white highlights; the rock warm grey with the spring water bright cyan and white ripple "
   "rings; the small pool pale blue; the three bamboo shoots pale yellow-green with warm brown wrapping "
   "sheaths; the wooden ladle and the small bucket warm tan; the bamboo fence and the footpath pale warm tan; "
   "the child's shirt warm yellow-orange, the neckerchief pure bright red, hair dark brown; the mushrooms warm "
   "red caps with cream stems; the grass tufts yellow-green; the distant bamboo clumps soft green; the birds "
   "pale grey-blue.",

 '3_嘉陵江边的一壶白开水':
   " Colour scheme: soft warm morning light; the sky pale blue with two thin pale mist bands; the wide river in "
   "muted blue-grey with white ripple lines; the two wooden boats warm brown with pale cream masts; the old "
   "street stone pavement warm light grey with paler paving-joint lines; the roof tiles dark slate grey with "
   "warm brown eave beams; the three hanging lanterns pure bright red with a warm yellow glow; the clay "
   "rain-water jar warm brown with a dark lid and a bamboo dipper in pale yellow; the bamboo chair and the long "
   "wooden bench warm tan; the cat pale orange with a white chest; the wooden table warm brown; the plain glass "
   "of water almost colourless pale blue with a white highlight swirl; the sugary drink in warm caramel brown "
   "with a light blue drinking straw; the sixteen sugar lumps pure white with pale grey edges; the "
   "long-spouted kettle warm copper-orange with a pale cream steam line; the white porcelain pot clean white "
   "with a pale blue rim; the child's shirt pale cream, the neckerchief pure bright red, hair dark brown; the "
   "bamboo-woven water bottle hanging on the wall pale yellow; the far tree deep green; the distant hills soft "
   "blue-grey; the birds pale grey-blue.",

 '4_城里的校园球赛_每人的水壶':
   " Colour scheme: bright cheerful daylight; the sky clear light blue with three white clouds and a big warm "
   "yellow sun; the running track brick red with white lane lines; the basketball court surface fresh green with "
   "white court lines; the hoop pole pale grey with a white backboard, a red rim and a white net; the basketball "
   "warm orange with darker seam lines; the two teaching buildings pale cream with light grey roofs and pale "
   "blue window glass; the corridor drinking-water station pale blue-white with three grey taps; the flower bed "
   "green with red and yellow flowers; the watering can bright blue; the plastic bucket pale green with rain "
   "water in light blue; the trees deep green and the bamboo clump yellow-green; the flag pole pale grey on a "
   "light grey platform; the mountain ridge soft blue-green; the children's sports shirts in bright yellow, sky "
   "blue and white with pure bright red neckerchiefs and navy blue shorts; their water bottles in bright red, "
   "bright blue and bright green with white stripe patterns and grey measuring marks; the towel white with a "
   "blue stripe; the sweat drops pale cyan; the referee teacher's shirt white with black shorts and a small "
   "silver whistle; the grass bright green; the birds pale grey-blue.",

 '5_我家的水位线':
   " Colour scheme: warm cosy indoor light; the wooden table, stool and chair warm honey brown; the wall pale "
   "cream; the child's three cups left to right in pale mint green, pale blue and light yellow with the water "
   "inside in light blue getting deeper to the right; the family's cups: the enamel mug clean white with a red "
   "rim stripe, the plain glass pale blue water, the vacuum thermos flask deep green with a silver lid, the "
   "handled mug warm orange; the small pouring pot warm copper; the stove dark grey with the metal kettle in "
   "warm copper and a pale cream curly steam line; the wall shelf pale wood with white bowls and dishes; the "
   "window frame pale wood with a short pale blue curtain; through the window the courtyard with a warm brown "
   "clay rain-water jar, a clothesline with pale blue and white clothes, a deep green pomelo tree with three "
   "pale yellow-green round fruit, a low pale grey wall and soft blue-grey distant hills; the fruit basket warm "
   "tan with three red and yellow fruit; the cat pale orange curled up beside the jar; the potted plant deep "
   "green in a warm terracotta pot; the grandparent's clothing pale grey-blue; the child's shirt warm yellow, "
   "hair dark brown; the sky pale blue with two white clouds.",
}


def colorize(path, strength=0.6):
    name = os.path.basename(path)
    key = name.rsplit('.', 1)[0]
    prompt = BASE_INSTR + COLORS.get(key, '')
    b64 = base64.b64encode(open(path, 'rb').read()).decode()
    for a in range(3):
        try:
            r = requests.post(f"{BASE}/images/generations",
                headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
                json={"model": MODEL, "prompt": prompt, "size": "2048x2048", "n": 1, "watermark": False,
                      "image": f"data:image/jpeg;base64,{b64}", "strength": strength}, 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}  s={strength}  {len(img)//1024}KB", flush=True)
                return p
            print(f"ERR {name} #{a}: {json.dumps(d, ensure_ascii=False)[:220]}", flush=True)
        except Exception as e:
            print(f"EXC {name} #{a}: {e}", flush=True)
        time.sleep(6)
    return None


if __name__ == '__main__':
    strength = float(sys.argv[1]) if len(sys.argv) > 1 else 0.6
    names = sys.argv[2:]
    files = sorted(glob.glob(os.path.join(SRC, '*.jpg')))
    if names:
        files = [f for f in files if any(n in os.path.basename(f) for n in names)]
    with ThreadPoolExecutor(max_workers=5) as ex:
        print(list(ex.map(lambda f: colorize(f, strength), files)))