#!/usr/bin/env python3
"""
Wan/Qwen 图像生成 + 编辑 Demo

该系列模型实际走 /v1/chat/completions 接口，
返回 choices[].message.content 数组里含 {image: url, type: image}

用法：
    NEWAPI_BASE_URL=https://your-gateway.example.com/v1 \
    NEWAPI_API_KEY=sk-xxx \
    python3 demo_wan_qwen_image.py               # 文生图

    NEWAPI_BASE_URL=https://your-gateway.example.com/v1 \
    NEWAPI_API_KEY=sk-xxx \
    python3 demo_wan_qwen_image.py --edit        # 图片编辑

环境变量：
    NEWAPI_BASE_URL     必填，网关地址（OpenAI 兼容 /v1 前缀）
    NEWAPI_API_KEY      必填，网关 API Key
    NEWAPI_MODEL        默认 wan2.7-image
    NEWAPI_PROMPT       提示词
    NEWAPI_INPUT_IMAGE  编辑模式的原图公网 URL
    NEWAPI_OUTPUT_FILE  输出文件名，默认 wan_qwen_demo.png
"""

import json
import os
import sys
import urllib.error
import urllib.request


BASE_URL = os.environ.get("NEWAPI_BASE_URL", "")
API_KEY = os.environ.get("NEWAPI_API_KEY", "")
MODE = "edit" if ("--edit" in sys.argv or os.environ.get("MODE") == "edit") else "gen"
MODEL = os.environ.get("NEWAPI_MODEL", "wan2.7-image")

DEFAULT_PROMPT = (
    "保留原图所有内容，把猫的毛色改成黑色"
    if MODE == "edit"
    else "一只橘猫坐在窗台上看夕阳，电影感光影，超写实摄影"
)
PROMPT = os.environ.get("NEWAPI_PROMPT", DEFAULT_PROMPT)
INPUT_IMAGE = os.environ.get(
    "NEWAPI_INPUT_IMAGE",
    "",  # 编辑模式需要传入公网可访问的图片 URL
)
OUTPUT_FILE = os.environ.get("NEWAPI_OUTPUT_FILE", "wan_qwen_demo.png")


def build_content():
    if MODE == "edit":
        return [
            {"type": "image_url", "image_url": {"url": INPUT_IMAGE}},
            {"type": "text", "text": PROMPT},
        ]
    return [{"type": "text", "text": PROMPT}]


def call_api():
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": build_content()}],
    }
    req = urllib.request.Request(
        f"{BASE_URL}/chat/completions",
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        return json.loads(resp.read().decode("utf-8"))


def extract_image_urls(resp):
    urls = []
    for choice in resp.get("choices", []):
        content = choice.get("message", {}).get("content")
        if isinstance(content, list):
            for item in content:
                if isinstance(item, dict) and item.get("image"):
                    urls.append(item["image"])
    return urls


def download(url: str, path: str):
    with urllib.request.urlopen(url, timeout=180) as r:
        with open(path, "wb") as fp:
            fp.write(r.read())


def main():
    if not API_KEY or not BASE_URL:
        print("请先设置环境变量 NEWAPI_BASE_URL 和 NEWAPI_API_KEY", file=sys.stderr)
        sys.exit(1)

    print(f"模式: {'图片编辑' if MODE == 'edit' else '文生图'}")
    print(f"模型: {MODEL}")
    print(f"Prompt: {PROMPT}")
    if MODE == "edit":
        print(f"原图: {INPUT_IMAGE}")
    print()

    try:
        resp = call_api()
    except urllib.error.HTTPError as e:
        print(f"HTTP {e.code}", file=sys.stderr)
        print(e.read().decode("utf-8", errors="replace"), file=sys.stderr)
        sys.exit(1)

    urls = extract_image_urls(resp)
    if not urls:
        print("响应无图片，原始数据：")
        print(json.dumps(resp, ensure_ascii=False, indent=2))
        sys.exit(1)

    base, ext = os.path.splitext(OUTPUT_FILE)
    for idx, url in enumerate(urls):
        out = OUTPUT_FILE if len(urls) == 1 else f"{base}_{idx + 1}{ext}"
        download(url, out)
        print(f"[{idx + 1}/{len(urls)}] ✅ {out}")
        print(f"     URL: {url[:120]}...")


if __name__ == "__main__":
    main()