# -*- coding: utf-8 -*-
"""横版 A4 · 第2步：配色图（i2i 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/horizontal/lineart'
OUT = '/data/www/files/waterHealthArt/horizontal/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. All water and drinking water must be painted in clean bright "
              "cyan-blue tones. ")

COLORS = {
 '1_一起喝够八杯水':
   " Colour scheme: clear light blue sky with three white clouds and a warm yellow sun; the plaza pavement and "
   "benches warm light beige stone with slightly darker joint lines; the tall stone monument warm beige with a "
   "stepped base and small pale grey lamp posts; the fountain basin pale grey stone with BRIGHT CYAN-BLUE water, "
   "a cyan water jet and white splash droplets; the big glass tumbler left almost colourless with pale blue "
   "water inside; the eight small cups pale cream with pale blue water; the wide river bright cyan-blue with "
   "white ripple lines; the small wooden boat warm brown; the mountain soft blue-green with darker green pines; "
   "the grass fresh green; the three children's shirts in sky blue, warm yellow and white, their neckerchiefs "
   "pure bright red, hair dark brown; the fruit on the plate warm orange and red.",

 '2_院坝里的一杯温开水':
   " Colour scheme: soft blue sky with two white clouds and a warm yellow sun; the courtyard ground warm pale "
   "brown sand, the courtyard wall warm cream with a warm brown wooden gate; the wall tap pale grey with a "
   "BRIGHT CYAN-BLUE straight stream of water and white splash droplets hitting the water in the wooden bucket; "
   "the bucket warm brown, the stone basin warm grey; the brick stove warm terracotta red-brown with a copper "
   "kettle on top and a pale cream curly steam line; the firewood warm brown; the low bench warm tan with a "
   "clean WHITE porcelain cooling jar with a pale blue rim standing on it; the pomelo tree deep green with pale "
   "yellow-green round fruit; the big clay rain-water jar warm brown with a bamboo dipper in pale yellow; the "
   "children's shirts in warm yellow and sky blue, their neckerchiefs pure bright red, hair dark brown; the "
   "glasses of water pale blue with a pale cream wisp of steam; the distant hill soft blue-green.",

 '3_白开水最好喝':
   " Colour scheme: soft warm daylight; the sky pale blue with two thin mist bands; the river muted blue-grey "
   "with white ripple lines; the wooden boats warm brown; the street pavement warm light grey stone with paler "
   "joint lines; the roof tiles dark slate grey with warm brown eave beams and three pure bright red round "
   "lanterns with a warm yellow glow; the wooden tables and bench warm brown; the tray of about twenty sugar "
   "lumps PURE WHITE with pale grey cube edges, piled into a small white hill; the tipped transparent bottle "
   "left almost colourless with only a grey edge line, the drinking straw light blue; the porcelain cooling jar "
   "clean white with a pale blue rim; the plate white with a red and a yellow fruit; the cat pale orange with a "
   "white chest; the two children's glasses of water almost colourless with pale blue and white sparkle shapes "
   "around them; their shirts pale cream and pale sky blue, their neckerchiefs pure bright red, hair dark brown; "
   "the big clay rain-water jar warm brown; the distant hills soft blue-grey; the far bank tree deep green.",

 '4_接一杯_喝一口_关紧它':
   " Colour scheme: bright cheerful daylight; the corridor wall and floor warm pale cream with long pale grey "
   "perspective lines; the corridor railing pale grey; the three wall taps pale grey with BRIGHT CYAN-BLUE water "
   "running into the bottle at the left tap with white spray droplets and two pale blue ripple rings, the middle "
   "cup pale cream with pale blue water, and the right tap completely dry with no water; the children's water "
   "bottles in bright red, bright blue and bright green with white stripe patterns and plain grey tick marks; "
   "their shirts in bright yellow, sky blue and white, shorts navy blue, neckerchiefs pure bright red; the two "
   "school teaching buildings pale cream with light grey roofs and pale blue window glass; the big plastic "
   "bucket pale green with cyan-blue rain water; the flower bed green with red and yellow flowers; the running "
   "track brick red with white lane lines; the basketball hoop pole pale grey with a white backboard and a red "
   "rim; the trees deep green, the bamboo clump yellow-green; the mountain ridge soft blue-green; the sky clear "
   "light blue with three white clouds and a warm yellow sun.",

 '5_全家一起_喝够每一杯':
   " Colour scheme: warm cosy indoor light; the wall pale cream, the long wooden table and chairs warm honey "
   "brown; the five cups left to right: a pale mint green bowl, a pale blue tall bottle, a plain glass with pale "
   "blue water, a warm orange handled mug, and a clean WHITE enamel mug with a bright red rim stripe, each with "
   "a clearly different pale-blue water level; the big porcelain cooling jar at the centre clean white with a "
   "pale blue rim; the brick stove warm terracotta with a copper kettle and a pale cream curly steam line; the "
   "wall shelf pale wood with white bowls and dishes; the window frame pale wood; 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 fruit, a low pale grey wall and soft blue-grey distant hills; "
   "the fruit basket warm tan with a red and a yellow fruit; the potted plant deep green in a warm terracotta "
   "pot; the child's shirt warm yellow, the adult's shirt soft blue, the grandparent's clothing pale grey-blue, "
   "hair dark brown; two white clouds in the pale blue sky.",
}


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": "2896x2048", "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)))