# LIB · 即梦4.0 (Seedream 4.0) 适配器
# D180 · 2026-07-10
"""火山方舟 ARK API → Seedream 4.0 图像生成"""
import json
import os
import urllib.request
import base64

ARK_KEY = "ark-ddeba9f4-8c5a-449e-b549-9c29ec1e6f8c-a39ea"
ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3"
SEEDREAM_MODEL = "doubao-seedream-4-0-250828"

def _api_call(endpoint, payload):
    """通用 ARK API 调用"""
    req = urllib.request.Request(
        f"{ARK_BASE}{endpoint}",
        data=json.dumps(payload).encode(),
        headers={
            "Authorization": f"Bearer {ARK_KEY}",
            "Content-Type": "application/json"
        },
        method="POST"
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        err_body = e.read().decode()
        return {"error": str(e.code), "detail": err_body}

def seedream_t2i(prompt, size="1440x2560", n=1):
    """文生图"""
    payload = {
        "model": SEEDREAM_MODEL,
        "prompt": prompt,
        "size": size,
        "n": n
    }
    return _api_call("/images/generations", payload)

def seedream_i2i(prompt, image_path, size="1440x2560", n=1):
    """图生图"""
    with open(image_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()
    payload = {
        "model": SEEDREAM_MODEL,
        "prompt": prompt,
        "image": f"data:image/jpeg;base64,{img_b64}",
        "size": size,
        "n": n
    }
    return _api_call("/images/generations", payload)

def seedream_multi(prompt, image_paths, size="1440x2560", n=1):
    """多图融合"""
    imgs = []
    for p in image_paths:
        with open(p, "rb") as f:
            fmt = "png" if p.endswith(".png") else "jpeg"
            imgs.append(f"data:image/{fmt};base64,{base64.b64encode(f.read()).decode()}")
    payload = {
        "model": SEEDREAM_MODEL,
        "prompt": prompt,
        "image": imgs,
        "size": size,
        "n": n
    }
    return _api_call("/images/generations", payload)

def seedream_download(result, output_path):
    """从 API 结果下载图片"""
    if "error" in result:
        return None
    url = result.get("data", [{}])[0].get("url")
    if not url:
        return None
    urllib.request.urlretrieve(url, output_path)
    return output_path

# ====== CLI ======
if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python seedream.py t2i <prompt> [output]")
        print("       python seedream.py i2i <prompt> <image_path> [output]")
        print("       python seedream.py multi <prompt> <img1> <img2> [output]")
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == "t2i":
        prompt = sys.argv[2]
        out = sys.argv[3] if len(sys.argv) > 3 else "seedream_output.jpg"
        r = seedream_t2i(prompt)
        p = seedream_download(r, out)
        print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")
    elif cmd == "i2i":
        prompt, img, out = sys.argv[2], sys.argv[3], (sys.argv[4] if len(sys.argv) > 4 else "seedream_i2i_output.jpg")
        r = seedream_i2i(prompt, img)
        p = seedream_download(r, out)
        print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")
    elif cmd == "multi":
        prompt, imgs, out = sys.argv[2], sys.argv[3:-1], sys.argv[-1]
        r = seedream_multi(prompt, imgs)
        p = seedream_download(r, out)
        print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")