#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "requests>=2.31.0",
#     "pillow>=10.0.0",
# ]
# ///
"""使用 Agnes AI（OpenAI 兼容网关）生成图片。
支持：文生图、图生图（参考图保持一致性）、URL/Base64 输出、下载保存、白底转透明。
"""

import argparse
import io
import json
import os
import sys
import time
import urllib.request

import requests

API_URL = "https://apihub.agnes-ai.com/v1/images/generations"
DEFAULT_MODEL = "agnes-image-2.1-flash"


def call_api(prompt, model, size, api_key, images=None, want_b64=False, retries=3):
    """调用 Agnes 图像生成接口，返回解析后的 JSON。
    对瞬时网络/SSL 抖动自动重试，重试间隔递增。
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    extra = {"response_format": "b64_json" if want_b64 else "url"}
    if images:
        # 图生图：参考图放在 extra_body.image，支持公网 URL 或 data:image Base64
        extra["image"] = images
    payload = {
        "model": model,
        "prompt": prompt,
        "size": size,
        "extra_body": extra,
    }
    # 文生图需要 Base64 输出时，按文档使用顶层 return_base64
    if want_b64 and not images:
        payload["return_base64"] = True

    last_err = None
    for attempt in range(1, retries + 1):
        try:
            resp = requests.post(API_URL, headers=headers, json=payload, timeout=180)
            resp.raise_for_status()
            return resp.json()
        except (requests.exceptions.SSLError,
                requests.exceptions.ConnectionError,
                requests.exceptions.Timeout) as e:
            # 瞬时网络抖动：重试
            last_err = e
            if attempt < retries:
                time.sleep(attempt * 2)
                continue
            raise
    raise last_err


def fetch_bytes(url):
    """下载 URL 内容为字节。"""
    with urllib.request.urlopen(url, timeout=180) as r:
        return r.read()


def make_transparent(png_bytes, thresh=60, white_min=225):
    """把纯白背景变透明（保留内部白色）。
    从图像四边多个种子点做 flood fill，只清除与边缘连通的白色区域。
    """
    from PIL import Image, ImageDraw

    img = Image.open(io.BytesIO(png_bytes)).convert("RGBA")
    w, h = img.size
    seeds = []
    steps = 12
    for i in range(steps + 1):
        t = int(i * (w - 1) / steps)
        seeds += [(t, 0), (t, h - 1)]
        s = int(i * (h - 1) / steps)
        seeds += [(0, s), (w - 1, s)]
    for sx, sy in seeds:
        r, g, b, a = img.getpixel((sx, sy))
        if a > 0 and r >= white_min and g >= white_min and b >= white_min:
            ImageDraw.floodfill(img, (sx, sy), (255, 255, 255, 0), thresh=thresh)
    out = io.BytesIO()
    img.save(out, "PNG")
    return out.getvalue()


def main():
    parser = argparse.ArgumentParser(description="使用 Agnes AI 生成图片。")
    parser.add_argument("--prompt", required=True, help="图像生成的文本描述（必选）")
    parser.add_argument("--model", default=DEFAULT_MODEL, help=f"模型 ID，默认 {DEFAULT_MODEL}")
    parser.add_argument("--size", default="1024x1024", help="图像尺寸，如 1024x1024、1024x768")
    parser.add_argument("--image", action="append", default=None,
                        help="参考图 URL 或 data:image Base64（图生图，可多次指定）")
    parser.add_argument("--api-key", help="Agnes API Key（也可用 AGNES_API_KEY 环境变量）")
    parser.add_argument("--save", help="下载并保存到本地路径，输出 SAVED:<path>")
    parser.add_argument("--transparent", action="store_true",
                        help="配合 --save：把纯白背景转为透明 PNG")
    parser.add_argument("--thresh", type=int, default=60, help="透明化白色容差，默认 60")

    args = parser.parse_args()

    api_key = args.api_key or os.environ.get("AGNES_API_KEY")
    if not api_key:
        print("ERROR: 缺少 API Key，请用 --api-key 或设置 AGNES_API_KEY 环境变量。")
        sys.exit(1)

    # 需要保存且要透明时，直接请求 URL 再下载处理（Pillow 处理本地字节）
    want_b64 = False
    try:
        result = call_api(
            prompt=args.prompt,
            model=args.model,
            size=args.size,
            api_key=api_key,
            images=args.image,
            want_b64=want_b64,
        )
    except requests.exceptions.RequestException as e:
        print(f"ERROR: API 请求失败: {e}")
        if getattr(e, "response", None) is not None:
            print(f"Response body: {e.response.text[:500]}")
        sys.exit(1)

    data = result.get("data") or []
    if not data:
        print(f"ERROR: 返回中无图像数据。完整响应: {json.dumps(result)[:500]}")
        sys.exit(1)

    item = data[0]
    url = item.get("url")
    b64 = item.get("b64_json")

    if not args.save:
        # 仅输出链接（对齐 seedream 约定）
        if url:
            print(f"MEDIA_URL: {url}")
        elif b64:
            print("ERROR: 收到 Base64 但未指定 --save，无法直接展示。请加 --save 保存。")
            sys.exit(1)
        else:
            print(f"ERROR: 无 url 也无 b64_json。响应: {json.dumps(result)[:500]}")
            sys.exit(1)
        return

    # 保存到本地
    try:
        if b64:
            import base64
            img_bytes = base64.b64decode(b64)
        elif url:
            img_bytes = fetch_bytes(url)
        else:
            print("ERROR: 无可保存的图像数据。")
            sys.exit(1)

        if args.transparent:
            try:
                img_bytes = make_transparent(img_bytes, thresh=args.thresh)
            except ImportError:
                print("ERROR: --transparent 需要 Pillow，请先 `pip install pillow`。")
                sys.exit(1)

        os.makedirs(os.path.dirname(os.path.abspath(args.save)), exist_ok=True)
        with open(args.save, "wb") as f:
            f.write(img_bytes)
        # 即使保存到本地，也输出原始图片链接（便于直接引用/展示）
        if url:
            print(f"MEDIA_URL: {url}")
        print(f"SAVED: {args.save}")
    except Exception as e:
        print(f"ERROR: 保存失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()