# -*- coding: utf-8 -*-
"""第2步：由水彩插画反推线条图（保构图、去色去阴影）—— 与 pandaRead 手抄报同一套做法
strength=0.6，prompt 强调纯黑白描、同形状同位置、无灰阶无排线无纸纹"""
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/final/illustration'
OUT = '/data/www/files/waterHealthArt/final/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. 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 outlines confident and hand-drawn with slight thickness variation. "
              "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)))

    # 中性问法复核主题（不给任何提示，看观众自己读成什么）
    from concurrent.futures import ThreadPoolExecutor as TPE
    KEY, M2 = CFG['api_key'], CFG['models']['text_vision']

    def name_of(p):
        b64 = base64.b64encode(open(p, 'rb').read()).decode()
        q = "请给这张儿童画起一个作品名（不超过10个字），并另起一行用一句话说明它想宣传的健康习惯。只回答，不要解释。"
        msgs = [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
            {"type": "text", "text": q}]}]
        r = requests.post(f"{BASE}/chat/completions", headers={"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"},
                          json={"model": M2, "messages": msgs, "max_tokens": 200}, timeout=180).json()
        return f"===== {os.path.basename(p)}\n{r['choices'][0]['message']['content'].strip()}\n"

    print("\n--- 中性问法复核 ---", flush=True)
    with TPE(max_workers=3) as ex:
        for o in ex.map(name_of, files):
            print(o, flush=True)
