# -*- coding: utf-8 -*-
"""final2：由 v6 插画反推线条图（保构图、去色去阴影）"""
import base64, glob, json, os, 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/final2/illustration'
OUT = '/data/www/files/waterHealthArt/final2/lineart'
os.makedirs(OUT, exist_ok=True)

LINE_INSTR = ("Convert this watercolour illustration into a PURE BLACK-AND-WHITE LINE DRAWING for children to trace "
              "and colour. Keep EXACTLY the same shapes, positions, proportions, element count and composition — do "
              "not add, remove, move or redraw anything. Keep the outlines confident and hand-drawn with clearly "
              "VARIED stroke thickness, some broken and dry-brush line ends. Remove ALL colour, ALL shading, ALL "
              "grey tones, ALL watercolour texture, ALL paper grain and ALL soft gradients — uniform clean black "
              "outlines on a pure white background, every enclosed area left pure white. NO cross-hatching, NO "
              "stippling, NO tonal patches, NO grey fill. Keep the drawn finger lines on the hands. NO text "
              "anywhere: no letters, no numbers, no characters, no labels, no signature, no watermark. ")


def reline(path, strength=0.6):
    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": LINE_INSTR, "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)[:200]}", flush=True)
        except Exception as e:
            print(f"EXC {name} #{a}: {e}", flush=True)
        time.sleep(6)
    return None


if __name__ == '__main__':
    files = sorted(glob.glob(os.path.join(SRC, '*.jpg')))
    with ThreadPoolExecutor(max_workers=3) as ex:
        print(list(ex.map(reline, files)))
