import base64
import getpass
import json
from pathlib import Path
from urllib.error import HTTPError
from urllib.request import Request, urlopen
import uuid
import webbrowser

URL = 'https://yyds.seqnode.com/v1/images/edits'
PAYLOAD = {
    "model": "grok-imagine-image-2.0",
    "prompt": "保留参考图的主体布局，将背景改为浅蓝色，增加柔和光线。",
    "images": [
        {
            "image_url": "https://yyds.seqnode.com/api-docs/parameter-test-fixtures/1.png"
        }
    ],
    "aspect_ratio": "1:1",
    "resolution": "1k",
    "quality": "low",
    "n": 1,
    "response_format": "b64_json"
}


def main():
    api_key = getpass.getpass("输入 API Key：").strip()
    if not api_key:
        raise ValueError("请输入有效的 API Key。")
    headers = {"Authorization": "Bearer " + api_key}
    headers["Content-Type"] = "application/json; charset=utf-8"
    headers["User-Agent"] = "YM2API-Quickstart/1.0"
    request = Request(URL, data=json.dumps(PAYLOAD, ensure_ascii=False).encode("utf-8"), headers=headers)
    with urlopen(request, timeout=300) as response:
        result = json.load(response)
    encoded = next((item["b64_json"] for item in result.get("data", []) if item.get("b64_json")), None)
    if not encoded:
        raise ValueError("响应没有最终图片。请检查响应和使用记录，再决定是否重试。")
    data = base64.b64decode(encoded, validate=True)
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        extension = ".png"
    elif data.startswith(b"\xff\xd8\xff"):
        extension = ".jpg"
    elif len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        extension = ".webp"
    else:
        raise ValueError("返回数据不是可识别的 PNG、JPEG 或 WebP。")
    path = Path.cwd() / ("grok-image-" + uuid.uuid4().hex + extension)
    with path.open("xb") as output:
        output.write(data)
    print(f"图片已保存：{path}")
    if not webbrowser.open(path.as_uri()):
        print("请到以上路径手动打开图片。")


if __name__ == "__main__":
    try:
        main()
    except HTTPError as error:
        print(f"HTTP {error.code}：{error.read(2000).decode('utf-8', errors='replace')}")
        raise SystemExit(1)
    except Exception as error:
        print(f"运行失败：{error}")
        raise SystemExit(1)