# encoding: utf-8
import uuid
import time
import requests
import base64
import json
from auth_util import gen_sign_headers

# 请替换APP_ID、APP_KEY
APP_ID = '3038561747'
APP_KEY = 'eMDGaCyRSRnPoQVb'
URI1 = '/api/v1/task_submit'    #文生图的给任务ID的
URI2 = '/api/v1/task_progress'  #文生图查询图片链接的
DOMAIN = 'api-ai.vivo.com.cn'
METHOD1 = 'POST'
METHOD2 = 'GET'
#通用v3 = '55c682d5eeca50d4806fd1cba3628781'
#通用v4 = '8fe3d641be3e589dad231dc6c3b1429a'
mode_list =['55c682d5eeca50d4806fd1cba3628781','8fe3d641be3e589dad231dc6c3b1429a']
def draw_vivogpt(x,y,z,m):

    task = submit(x, y, z, m)
    time.sleep(3)  # 延迟5秒
    images_url = progress(task)
    while images_url == 0:
        time.sleep(2)  # 延迟2秒
        images_url = progress(task)
        print(images_url)
    return images_url
def submit(x,y,z,m):
   mode =mode_list[m]
   params = {}
   data = {
        'height': x,
        'width': y,
        'prompt': z,
        'styleConfig': mode
   }


   headers = gen_sign_headers(APP_ID, APP_KEY, METHOD1, URI1, params)
   headers['Content-Type'] = 'application/json'

   url2 = 'http://{}{}'.format(DOMAIN, URI1)
   response = requests.post(url2, data=json.dumps(data), headers=headers)
   if response.status_code == 200:
       response_data = response.json()
       # 从响应数据中获取task_id
       task_id = response_data.get('result', {}).get('task_id')
       if task_id:
           print(task_id)
           return task_id
       else:
           print("Task ID not found in response.")
           return None
   else:
       print(f"Error: {response.status_code}, {response.text}")
       return None
def progress(task_id):
   params = {

       'task_id': task_id  #'35c81976cd9f559783139a489ec2bef1'
   }
   headers = gen_sign_headers(APP_ID, APP_KEY, METHOD2, URI2, params)

   uri_params = ''
   for key, value in params.items():
       if key is not None and value is not None:
           uri_params += f"{key}={value}&"

   uri_params = uri_params[:-1]

   url = 'http://{}{}?{}'.format(DOMAIN, URI2, uri_params)

   response = requests.get(url, headers=headers)
   if response.status_code == 200:
       response_data = response.json()
       images_url = response_data.get('result', {}).get('images_url', [])
       if images_url==[]:
          print('请等待')
          return 0
       else:
           return images_url[0]

   else:
       print(response.status_code, response.text)


if __name__ == '__main__':
    draw_vivogpt(x,y,z,m)