#!/usr/bin/env python3
"""调用 grok-imagine-image-lite 生成图片的示例脚本"""

import argparse
import base64
import datetime
import json
import sys
import urllib.request
import urllib.error
from pathlib import Path

API_URL = "https://api.slomerex.xyz/v1/images/generations"
API_KEY = "sk-f8c4T7gx9Putj039gca11sstL8rTYXr1rSDJmGf6Cg0IhOw8"
MODEL = "grok-imagine-image-lite"

# 图片描述提示词，直接改这里即可
PROMPT = """A majestic white-and-gold dragon soaring through a vast sky filled with drifting clouds at golden hour. Its scales glisten under the warm sunset light, and its wings cast soft shadows over a distant misty mountain range below. Cinematic composition, photorealistic textures, volumetric lighting, 8K, extremely detailed."""


def generate_image(prompt: str, size: str = "1024x1024", output: str = "output.png"):
    payload = {
        "model": MODEL,
        "prompt": prompt,
        "size": size,
        "n": 1,
    }

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
    }

    req = urllib.request.Request(
        API_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        print(f"HTTP {e.code} 错误: {body}", file=sys.stderr)
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"网络错误: {e}", file=sys.stderr)
        sys.exit(1)

    if "data" not in data:
        print(f"响应中没有 data 字段: {json.dumps(data, ensure_ascii=False)}", file=sys.stderr)
        sys.exit(1)

    out_path = Path(output)
    for i, item in enumerate(data["data"]):
        # 优先取 b64_json，否则取 url
        if item.get("b64_json"):
            raw = base64.b64decode(item["b64_json"])
        elif item.get("url"):
            # 图片 URL 需要带浏览器 User-Agent，否则会被服务端拦截（403）
            img_req = urllib.request.Request(item["url"], headers={"User-Agent": headers["User-Agent"]})
            with urllib.request.urlopen(img_req, timeout=120) as img_resp:
                raw = img_resp.read()
        else:
            print(f"第 {i} 张图片没有 b64_json / url 字段", file=sys.stderr)
            continue

        # Grok 默认返回 JPEG，根据 mime_type 修正扩展名
        ext_map = {"image/jpeg": ".jpg", "image/png": ".png", "image/webp": ".webp"}
        mime = item.get("mime_type") or item.get("content_type")
        path = out_path
        if len(data["data"]) > 1:
            path = out_path.with_name(f"{out_path.stem}_{i}{out_path.suffix}")
        if mime and mime in ext_map and path.suffix.lower() != ext_map[mime]:
            path = path.with_suffix(ext_map[mime])
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(raw)
        print(f"已保存: {path} ({len(raw)} bytes)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用 grok-imagine-image-lite 生成图片")
    parser.add_argument("prompt", nargs="?", default=None, help="图片提示词（不传则使用脚本内硬编码的 PROMPT）")
    parser.add_argument("--size", default="1024x1024", help="图片尺寸，如 1024x1024")
    # 默认输出到项目根目录的 output/ 文件夹，按「模型名-时间戳.后缀名」命名
    parser.add_argument(
        "--output", "-o",
        default=f"output/{MODEL}-{datetime.datetime.now():%Y%m%d_%H%M%S}.png",
        help="输出文件路径（默认 output/ 目录下「模型名-时间戳.png」，实际后缀按返回的 mime_type 自动修正）",
    )
    args = parser.parse_args()

    generate_image(args.prompt or PROMPT, args.size, args.output)