#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""GLM-Image 文生图调用脚本（智谱 GLM-Image 模型）。

用法:
  python3 generate.py "提示词" [-s SIZE] [-o OUTPUT] [--open] [--json]

示例:
  python3 generate.py "一只可爱的小猫咪，坐在阳光明媚的窗台上" -s 1280x1280
  python3 generate.py "商业海报：新品上市" -s 1056x1568 -o poster.png --open

默认 size=1280x1280，输出到当前目录 glm-image-<时间戳>.png
API Key 优先读环境变量 GLM_API_KEY，否则用内置默认 key。
"""
import argparse
import json
import os
import sys
import time
import urllib.request
import urllib.error

API_ENDPOINT = "https://open.bigmodel.cn/api/paas/v4/images/generations"
DEFAULT_KEY = ""  # 服务器版不含内置 key，请设置环境变量 GLM_API_KEY

RECOMMENDED_SIZES = [
    "1280x1280", "1568x1056", "1056x1568",
    "1472x1088", "1088x1472", "1728x960", "960x1728",
]


def parse_size(size):
    """校验 size，返回 (w, h)。规则：512-2048，且为 32 的整数倍。"""
    try:
        w, h = size.lower().split("x")
        w, h = int(w), int(h)
    except ValueError:
        raise ValueError(f"size 格式错误：'{size}'，应为 WxH，如 1280x1280")
    for v, name in ((w, "宽"), (h, "高")):
        if v < 512 or v > 2048:
            raise ValueError(f"{name}={v} 不在 512-2048 范围内")
        if v % 32 != 0:
            raise ValueError(f"{name}={v} 不是 32 的整数倍")
    return w, h


def generate(prompt, size="1280x1280", api_key=None, timeout=120, watermark=True):
    """调用 GLM-Image 接口，返回图片 URL。

    watermark=False 关闭 AI 水印，需账号已在「个人中心-安全管理-去水印管理」签署免责声明。
    """
    parse_size(size)  # 校验
    key = api_key or os.environ.get("GLM_API_KEY") or DEFAULT_KEY
    payload = {
        "model": "glm-image",
        "prompt": prompt,
        "size": size,
        "watermark_enabled": watermark,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        API_ENDPOINT,
        data=data,
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        err = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {err}") from None
    except urllib.error.URLError as e:
        raise RuntimeError(f"网络错误: {e.reason}") from None

    obj = json.loads(body)
    if not obj.get("data"):
        raise RuntimeError(f"返回无 data 字段: {body}")
    url = obj["data"][0].get("url")
    if not url:
        raise RuntimeError(f"返回无 url: {body}")
    return url, obj


def download(url, output):
    """下载图片 URL 到 output，返回输出路径。"""
    with urllib.request.urlopen(url, timeout=120) as resp:
        content = resp.read()
    with open(output, "wb") as f:
        f.write(content)
    return output


def main():
    ap = argparse.ArgumentParser(description="GLM-Image 文生图")
    ap.add_argument("prompt", help="生成提示词（最多 1000 字符）")
    ap.add_argument("-s", "--size", default="1280x1280",
                    help=f"图片尺寸 WxH（默认 1280x1280）。推荐: {', '.join(RECOMMENDED_SIZES)}")
    ap.add_argument("-o", "--output", default=None,
                    help="输出文件路径（默认 glm-image-<时间戳>.png）")
    ap.add_argument("--open", action="store_true", help="生成后在 Finder 中打开")
    ap.add_argument("--json", action="store_true", help="打印完整 JSON 返回")
    ap.add_argument("--no-download", action="store_true", help="只返回 URL，不下载")
    ap.add_argument("--no-watermark", action="store_true",
                    help="关闭 AI 水印（需账号已在「个人中心-安全管理-去水印管理」签署免责声明）")
    args = ap.parse_args()

    if len(args.prompt) > 1000:
        sys.exit(f"错误: 提示词 {len(args.prompt)} 字符，超过 1000 上限")

    try:
        parse_size(args.size)
    except ValueError as e:
        sys.exit(f"错误: {e}")

    print(f"→ 调用 GLM-Image（size={args.size}, watermark={'off' if args.no_watermark else 'on'}）...", file=sys.stderr)
    t0 = time.time()
    try:
        url, obj = generate(args.prompt, args.size, watermark=not args.no_watermark)
    except RuntimeError as e:
        sys.exit(f"错误: {e}")
    elapsed = time.time() - t0
    print(f"✓ 生成成功（{elapsed:.1f}s）: {url}", file=sys.stderr)

    if args.json:
        print(json.dumps(obj, ensure_ascii=False, indent=2))

    if args.no_download:
        print(url)
        return

    output = args.output or f"glm-image-{int(time.time())}.png"
    try:
        download(url, output)
    except Exception as e:
        sys.exit(f"下载失败: {e}\n图片 URL: {url}")
    print(f"✓ 已保存: {output}", file=sys.stderr)
    print(output)

    if args.open:
        os.system(f'open "{output}"')


if __name__ == "__main__":
    main()