# -*- coding: utf-8 -*-
"""第2步：以线条图为参考生成配色图（图生图）。
用法: python3 gen_color5.py <strength> [主题名...]
不带主题名则处理全部5份。"""
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/pandaReadArt/lineart5'
OUT = '/data/www/files/pandaReadArt/color5'
os.makedirs(OUT, exist_ok=True)

COLOR = ("Create the watercolour COLOURED version of this exact line drawing. Keep exactly the same shapes, the "
         "same positions, the same proportions and the same simple composition — do not add or remove any "
         "element, do not redraw anything. Simply fill the existing outlined areas with clean soft flat "
         "watercolour washes: transparent, gentle, evenly applied, staying neatly inside the outlines with the "
         "black outlines still clearly visible on top. Cheerful bright childlike palette, white paper still "
         "showing in the empty background, subtle paper grain. NO shading, NO gradients, NO texture, NO text of "
         "any kind, no letters, no characters, no numbers, no signature, no watermark, no stamp. Flat simple "
         "colouring-book style result, very easy for a child to copy.")


def colorize(path, strength):
    name = os.path.basename(path)
    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": COLOR, "size": "2048x2048", "n": 1, "watermark": False,
                      "image": f"data:image/jpeg;base64,{b64}", "strength": strength}, timeout=240)
            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(5)
    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)))