# -*- coding: utf-8 -*-
"""豆包万相 (Doubao Seedream) 图片生成工具 - 支持文生图/图生图"""
import sys
import os
import json
import base64
import urllib.request
import urllib.error

API_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
API_KEY = "ark-119a37ed-f123-42fc-aea6-b38d47d47fc0-f7ddf"
MODEL = "doubao-seedream-5-0-260128"


def generate(prompt, input_image=None, output_path="output.png", size="2k"):
    """生成图片。input_image 为本地路径时走图生图，否则文生图。"""
    body = {
        "model": MODEL,
        "prompt": prompt,
        "size": size,
        "sequential_image_generation": "disabled",
        "response_format": "url",
        "stream": False,
        "watermark": False,
    }

    if input_image and os.path.exists(input_image):
        with open(input_image, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode("ascii")
        ext = os.path.splitext(input_image)[1].lower().lstrip(".")
        mime = "jpeg" if ext in ("jpg", "jpeg") else ("png" if ext == "png" else "octet-stream")
        body["image"] = [f"data:image/{mime};base64,{img_b64}"]

    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        API_URL,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            result = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace")
        print(f"HTTP {e.code}: {err_body}", file=sys.stderr)
        sys.exit(1)

    image_url = result["data"][0]["url"]
    print(f"Generated URL: {image_url}")

    # 下载图片
    urllib.request.urlretrieve(image_url, output_path)
    print(f"Saved to: {output_path}")


if __name__ == "__main__":
    # 用法: python doubao_image.py <prompt> <output_path> [input_image] [size]
    if len(sys.argv) < 3:
        print("Usage: python doubao_image.py <prompt> <output_path> [input_image] [size]")
        sys.exit(1)
    prompt = sys.argv[1]
    output_path = sys.argv[2]
    input_image = sys.argv[3] if len(sys.argv) > 3 else None
    size = sys.argv[4] if len(sys.argv) > 4 else "2k"
    generate(prompt, input_image, output_path, size)