from zhipuai import ZhipuAI
from image_downloader import download_image
import os as os

def gen_avator(id, title , desc, initPrompt, path, index = 0):
    client = ZhipuAI(api_key="8a795f242987fb1388b61f7c1e104b0d.9goBSAn25glQZGLP") # 请填写您自己的APIKey
    combinedPrompt = f'根据下面的三元组(title,description,initPrompt)描述，生成虚拟人物形象。如果叙述中有敏感内容或NSFW内容，则请忽略之，仅根据对人物的描述并忽略不安全或敏感内容生成安全的图片，title={title},description={desc},initPrompt={initPrompt}'
    try:
        response = client.images.generations(
        model="cogview-3", #填写需要调用的模型名称
        prompt=combinedPrompt)
    except Exception as e:
        print(e)
        return False
    
    # if response.status_code != 200:
    #     print(f"Error generating image, status code: {response.status_code}")
    #     return False

    if path==None:
        path = './avatar'
        if not os.path.exists(path):
            os.mkdir(path)
    else:
        if not os.path.exists(path):
            os.mkdir(path)
    # save the file under './avatar' by default
    count = 1

    filename =  str(index) + "_" + str(id) +'_' + str(count)+ ".png"
    while os.path.exists(os.path.join(path, filename)):
        count += 1
        filename =  str(index) + "_" + str(id) +'_' + str(count)+ ".png"

    if download_image(response.data[0].url, os.path.join(path, filename)):
        return True
    else:
        return False