"""Image generation API client (Doubao/Seedream)."""

import base64
import os
import urllib.request
from pathlib import Path

_RATIO_MAP = {
    "1:1": "1024x1024",
    "4:3": "1024x768",
    "3:4": "768x1024",
    "16:9": "1280x720",
    "9:16": "720x1280",
}


def generate_image(
    api_key: str,
    prompt: str,
    size: str | None = None,
    ratio: str | None = None,
    output: str | None = None,
    watermark: bool = True,
    b64: bool = False,
) -> str:
    """Generate image using Doubao/Seedream API.

    Returns absolute file path if saved, or URL string if output is None.
    """
    from openai import OpenAI

    resolved_size = size or (ratio and _RATIO_MAP.get(ratio)) or "2K"

    client = OpenAI(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=api_key,
    )

    resp = client.images.generate(
        model="doubao-seedream-5-0-260128",
        prompt=prompt,
        size=resolved_size,
        response_format="b64_json" if b64 else "url",
        extra_body={"watermark": watermark},
    )

    if b64:
        image_bytes = base64.b64decode(resp.data[0].b64_json)
        save_path = output or "generated_image.png"
        Path(save_path).write_bytes(image_bytes)
        return os.path.abspath(save_path)
    else:
        url = resp.data[0].url
        if output:
            urllib.request.urlretrieve(url, output)
            return os.path.abspath(output)
        return url