# LIB · Seedream 3/4/5 适配器
# D181 · 2026-07-10 · 默认模型: Seedream 3.0 (轻量,便宜)
"""火山方舟 ARK API → Seedream 系列 文生图/图生图"""
import json
import os
import urllib.request
import base64

from lib.secrets import get_ark_key
ARK_KEY = get_ark_key()
ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3"

# 模型配置
# 模型配置（ARK端点上实际可用的模型）
SEEDREAM_3 = "doubao-seedream-4-0-250828"       # Seedream 4.0 (最便宜可用)
SEEDREAM_4 = "doubao-seedream-4-0-250828"       # 同上
SEEDREAM_4_5 = "doubao-seedream-4-5-251128"     # Seedream 4.5
SEEDREAM_5 = "doubao-seedream-5-0-260128"       # Seedream 5.0
SEEDREAM_5_PRO = "doubao-seedream-5-0-pro-260628" # Seedream 5.0 Pro
SEEDREAM_MODEL = SEEDREAM_3                      # 默认模型
SEEDREAM_MODEL = SEEDREAM_3                        # 默认模型 (爸爸指定用便宜的)

# 尺寸约束
SIZES = {
    SEEDREAM_3: ["1024x1024", "1024x768", "768x1024"],
    SEEDREAM_4: ["1664x936", "936x1664", "1024x1024"],
}
DEFAULT_SIZE = "1024x1024"  # 即梦3默认方图

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=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
    """文生图
    prompt: 中英文提示词
    size:   "1024x1024" (即梦3) / "1664x936" (即梦4)
    model:  SEEDREAM_3 (默认) / SEEDREAM_4
    """
    # 即梦3不支持 n>1
    if model == SEEDREAM_3 and n > 1:
        import warnings
        warnings.warn("即梦3.0 不支持 n>1，已固定为 n=1")
        n = 1
    payload = {
        "model": model,
        "prompt": prompt,
        "size": size,
        "n": n
    }
    return _api_call("/images/generations", payload)

def seedream_i2i(prompt, image_path, size=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
    """图生图 · 即梦3.0 图生图 (seededit-3.0)"""
    if model == SEEDREAM_3:
        # 即梦3图生图用 seededit-3.0
        model = "seededit-3.0"
    with open(image_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()
    payload = {
        "model": model,
        "prompt": prompt,
        "image": f"data:image/jpeg;base64,{img_b64}",
        "size": size,
        "n": 1  # 图生图只支持单张
    }
    return _api_call("/images/generations", payload)

def seedream_multi(prompt, image_paths, size=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
    """多图融合 · 仅即梦4支持"""
    if model == SEEDREAM_3:
        raise RuntimeError("多图融合仅即梦4.0支持，请使用 model=SEEDREAM_4")
    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": 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] [--model 3|4] [--size WxH]")
        print("       python seedream.py i2i <prompt> <image_path> [output] [--model 3|4]")
        print("       python seedream.py multi <prompt> <img1> <img2> [output]")
        print(f"       Default model: {SEEDREAM_MODEL} (即梦3.0 · ¥0.2/张)")
        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')}")