import os
import sys
import requests
import json


def ai_picture_book_task_query(api_key: str, task_id: str):
    url = "https://qianfan.baidubce.com/v2/tools/ai_picture_book/query"
    headers = {
        "Authorization": "Bearer %s" % api_key,
        "X-Appbuilder-From": "openclaw",
        "Content-Type": "application/json"
    }
    task_ids = task_id.split(",")
    params = {
        "task_ids": task_ids,
    }
    response = requests.post(url, headers=headers, json=params)
    response.raise_for_status()
    result = response.json()
    datas = []
    if "code" in result:
        raise RuntimeError(result["detail"])
    if "errno" in result and result["errno"] != 0:
        raise RuntimeError(result["errmsg"])
    if "data" in result and len(result["data"]) > 0:
        for item in result["data"]:
            if item["result"]:
                bosUrl = item["result"].get("video_bos_url", "")
                del item["result"]
                item["video_bos_url"] = bosUrl
            datas.append(item)
    return datas


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python ai_picture_book_task_query.py <task_ids>")
        sys.exit(1)

    task_ids = sys.argv[1]

    api_key = os.getenv("BAIDU_API_KEY")
    if not api_key:
        print("Error: BAIDU_API_KEY must be set in environment.")
        sys.exit(1)
    try:
        results = ai_picture_book_task_query(api_key, task_ids)
        print(json.dumps(results, ensure_ascii=False, indent=2))
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)