# coding:utf-8
# author: 湖南涉外经济学院信息与机电工程学院
# date: 2022.09.09

from flask import Flask, render_template, request, jsonify
from flask import Blueprint, render_template
import requests
import json
import time
import urllib.request
import base64
from io import BytesIO
from PIL import Image

# 设置允许的文件格式
index_ernie_vilg = Blueprint("ernie_vilg", __name__)

global API_KEY, SECRET_KEY
API_KEY = ''   # 在https://console.bce.baidu.com/tools/ 去申请	ERNIE-ViLG 的API
SECRET_KEY = ''


def get_access_token():
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))


def get_result(taskId, access_token):
    url = f"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/getImg?access_token={access_token}"

    payload = json.dumps({
        "taskId": taskId
    })
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload).json()

    return response


# 接收页面传递参数
@index_ernie_vilg.route('/ernie_vilg', methods=['POST', 'GET'])  # 添加路由
def ernie_vilg():
    if request.method == 'POST':
        try:
            ernie_vilg_text = request.form.get('describe')
            ernie_vilg_num = request.form.get('selected_number')
            selected_style = request.form.get('selected_color')
            num = int(ernie_vilg_num)
            print(ernie_vilg_text, '\n=================\n', selected_style, '\n=================\n', request.form.get)
            file_list = ernie_vilg_gen(ernie_vilg_text, selected_style, num)
            # file_list = [f'static/images/target/{fff}' for fff in file_list]
            print(file_list)
            sourcefile0 = 'static/images/target/result_0.png'
            sourcefile1 = 'static/images/target/result_1.png'
            sourcefile2 = 'static/images/target/result_2.png'
            sourcefile3 = 'static/images/target/result_3.png'
            sourcefile4 = 'static/images/target/result_4.png'
            sourcefile5 = 'static/images/target/result_5.png'
            for filename in file_list:
                print(filename)
            return render_template('ernie_vilg.html', sourcefile1=sourcefile1, sourcefile2=sourcefile2,
                                   sourcefile3=sourcefile3, sourcefile4=sourcefile4, sourcefile5=sourcefile5,
                                   sourcefile0=sourcefile0, file_list=file_list)
        except Exception as e:
            print(e, "*" * 50)
            return render_template('ernie_vilg.html')
    return render_template('ernie_vilg.html')


# 生成
def ernie_vilg_gen(ernie_vilg_text, selected_style, num):
    # 发送HTTP请求
    print(f'你的API_KEY：{API_KEY}')
    print(f'你的API_KEY：{SECRET_KEY}')
    access_token = get_access_token()
    print(f'你的token：{access_token}')
    print("结果: selected_style=", selected_style)
    url = "https://aip.baidubce.com/rpc/2.0/ernievilg/v1/txt2img?access_token=" + access_token
    payload = json.dumps({
        "text": ernie_vilg_text,
        "resolution": "1024*1024",
        "style": selected_style,
        "num": num
    })
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload).json()
    print('###===###===###===###')
    print(response)
    print('###===###===###===###')
    taskId = response['data']['taskId']
    print(taskId)
    time.sleep(2)
    result = get_result(taskId, access_token)
    print(result)
    file_list = []
    while result['data']['imgUrls'] == []:
        time.sleep(2)
        result = get_result(taskId, access_token)
        print(result)
    for i in range(0, len(result['data']['imgUrls'])):
        url = result['data']['imgUrls'][i]['image']
        urllib.request.urlretrieve(url, 'static/images/target/' + f'result_{i}.png')
    return file_list