# -*- coding: utf-8 -*-
"""Генерация логотипа radiOLA через routerai.ru (OpenAI-совместимый, gpt-5.4-image-2)."""
import os, sys, json, base64, urllib.request, re

KEY = open(os.path.expanduser("~/.routerai_key")).read().strip()
BASE = "https://routerai.ru/api/v1"
OUT = r"C:\radiOLA\design\logos\gen"
os.makedirs(OUT, exist_ok=True)

def _data_uri(path):
    ext = os.path.splitext(path)[1].lstrip(".").lower() or "png"
    if ext == "jpg": ext = "jpeg"
    return f"data:image/{ext};base64," + base64.b64encode(open(path, "rb").read()).decode()

def call(model, prompt, ref_images=None):
    if ref_images:
        content = [{"type": "text", "text": prompt}]
        for p in ref_images:
            content.append({"type": "image_url", "image_url": {"url": _data_uri(p)}})
    else:
        content = prompt
    body = json.dumps({
        "model": model,
        "modalities": ["image", "text"],
        "messages": [{"role": "user", "content": content}],
    }).encode("utf-8")
    req = urllib.request.Request(BASE + "/chat/completions", data=body, headers={
        "Authorization": "Bearer " + KEY,
        "Content-Type": "application/json",
    })
    with urllib.request.urlopen(req, timeout=180) as r:
        return json.loads(r.read().decode("utf-8"))

def save_images(resp, stem):
    """Ищем картинки в ответе (разные возможные места) и сохраняем."""
    saved = []
    raw = json.dumps(resp)
    # 1) message.images[].image_url.url (data:image/...;base64,...)
    msg = (resp.get("choices") or [{}])[0].get("message", {}) or {}
    cands = []
    for im in (msg.get("images") or []):
        u = (im.get("image_url") or {}).get("url") or im.get("url")
        if u: cands.append(u)
    # 2) content может быть списком с image_url
    c = msg.get("content")
    if isinstance(c, list):
        for part in c:
            if isinstance(part, dict):
                u = (part.get("image_url") or {}).get("url")
                if u: cands.append(u)
    # 3) любые data:image в сыром json
    cands += re.findall(r'data:image/[^"\\]+;base64,[A-Za-z0-9+/=]+', raw)
    seen=set()
    for i, u in enumerate(cands):
        if u in seen: continue
        seen.add(u)
        m = re.search(r'base64,([A-Za-z0-9+/=]+)', u)
        if not m: continue
        p = os.path.join(OUT, f"{stem}_{i}.png")
        open(p, "wb").write(base64.b64decode(m.group(1)))
        saved.append(p)
    return saved

STYLE = ("App icon design, a single chunky glossy 3D cartoon letter R, vibrant emerald and "
    "lime green with bright glassy highlights and smooth rounded bevels, playful premium mascot "
    "style, centered on a deep dark teal-green background with a soft green glow, no text, "
    "square 1:1 composition, crisp high detail. ")

VARIANTS = {
    "v1_radio":  STYLE + "A small retro radio is built into the left stem of the R: round tuning "
        "knobs and a horizontal speaker grille, integrated cleanly into the letter.",
    "v2_waves":  STYLE + "Concentric broadcast signal arcs and a tiny antenna emit from the top-right "
        "of the R, like an on-air radio broadcast.",
    "v3_vinyl":  STYLE + "A glossy black vinyl record with a glowing green center label is tucked "
        "behind the lower-left of the R, hinting at music radio.",
    "v4_eq":     STYLE + "A row of glowing green equalizer sound bars of varying heights runs along "
        "the bottom, and the 3D R rises above them.",
    "v5_dial":   STYLE + "A retro radio tuning dial with a frequency scale and a glowing needle "
        "wraps around the base of the R.",
}

if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == "batch":
    model = "openai/gpt-5.4-image-2"
    for stem, prompt in VARIANTS.items():
        try:
            resp = call(model, prompt)
            s = save_images(resp, stem)
            print(stem, "->", s if s else "NO IMAGE", flush=True)
        except Exception as e:
            print(stem, "ERROR", e, flush=True)
    sys.exit(0)

# режим: python routerai.py refbatch <img>  — рендер монограммы в 3D, 4 стиля
if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == "refbatch":
    img = sys.argv[2]
    KEEP = ("Recreate the EXACT letter-R monogram shown in the reference image — keep its identical "
        "shape, proportions and thick ribbon construction, do not redesign the letterform. ")
    BG = " Centered on a dark teal-green background with a subtle green glow, app icon, no text, square 1:1, crisp high detail."
    PR = {
        "r3d_matte": KEEP + "Render it as a chunky 3D isometric extruded block letter with clean matte "
            "emerald-green faces and a darker green extruded side, soft simple shading, modern minimal 3D." + BG,
        "r3d_gloss": KEEP + "Render it as a glossy 3D isometric extruded letter, vibrant emerald and lime "
            "green with bright glassy highlights and smooth rounded bevels, premium look." + BG,
        "r3d_grad":  KEEP + "Render it as a 3D isometric extruded letter with a smooth lime-to-green gradient "
            "front face, a clean darker-green side, gentle top light, modern." + BG,
        "r3d_neon":  KEEP + "Render it as a sleek 3D isometric extruded letter in bright lime green with a "
            "soft neon glow rim, dark glossy side, futuristic clean." + BG,
    }
    for stem, prompt in PR.items():
        try:
            resp = call("openai/gpt-5.4-image-2", prompt, ref_images=[img])
            print(stem, "->", save_images(resp, stem) or "NO IMAGE", flush=True)
        except Exception as e:
            print(stem, "ERROR", e, flush=True)
    sys.exit(0)

# режим: python routerai.py ref <stem> "<prompt>" <img1> [img2 ...]
if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == "ref":
    stem = sys.argv[2]; prompt = sys.argv[3]; refs = sys.argv[4:]
    resp = call("openai/gpt-5.4-image-2", prompt, ref_images=refs)
    print("SAVED:", save_images(resp, stem))
    sys.exit(0)

if __name__ == "__main__":
    model = sys.argv[1] if len(sys.argv) > 1 else "openai/gpt-5.4-image-2"
    stem = sys.argv[2] if len(sys.argv) > 2 else "probe"
    prompt = sys.argv[3] if len(sys.argv) > 3 else (
        "App icon, a chunky glossy 3D cartoon letter 'R' in vibrant emerald and lime green "
        "with a bright glassy top highlight and a clean beveled extrusion. A small retro radio "
        "is built into the left stem of the R: round tuning knobs and horizontal speaker-grille "
        "lines, subtle. Centered on a dark teal-green background with a soft green glow. Bold, "
        "modern, playful, high detail, no text, square 1:1.")
    resp = call(model, prompt)
    # дамп структуры без гигантского base64
    skel = json.loads(re.sub(r'[A-Za-z0-9+/=]{200,}', '<BASE64>', json.dumps(resp)))
    print("STRUCTURE:", json.dumps(skel, ensure_ascii=False)[:1500])
    print("SAVED:", save_images(resp, stem))
