import asyncio
import json
import time
import requests

import re
import openai
from loguru import logger
from Helper import RedisCon, Helper

# key
api_key = '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'
# 组织标识
organization = 'org-cflsyp3YsvDovzOu3g8G9OFS'


# 回调接口
def gpt_callback(res_prompt, task):
    try:
        url = task['callback']
        if len(res_prompt):
            params = {
                'prompt': json.dumps(res_prompt),
                'user_id': task['user_id'],
                'task_id': task['task_id'],
                'prompt_img': task['prompt_img'],
                'type': task['type'],
            }
            res = requests.post(url, data=params)
        else:
            logger.error(f'指令任务不正确，{res_prompt}')
    except Exception as e:
        logger.error(e)


# 拆分指令
def split_prompt(response, task):
    try:
        text = response.choices[0].text
        if text:
            # 使用正则定义分隔符
            pattern = r'\n[1-9]{1,2}.*.'
            # 拆分字符串
            chemical_list = re.findall(pattern, text)
            for i in range(len(chemical_list)):
                # 匹配标题并替换
                chemical_list[i] = re.sub(r'\n[1-9]{1,3}.', '', chemical_list[i], count=0, flags=0)
            return chemical_list
    except Exception as e:
        logger.error(e)
        logger.error('gpt指令拆分失败，请查看原因')
        return []


# 获取模型列表
def get_module_list():
    header = {
        'Authorization': 'Bearer ' + api_key,
    }
    x = requests.get('https://api.openai.com/v1/models', headers=header)
    print(x.text)


def get_gpt_message(task):
    openai.organization = organization
    openai.api_key = api_key
    prompt_num = task['prompt_num']
    # 指令
    prompt = task['prompt']
    logger.info('开始访问open_ai')
    try:
        response = openai.Completion.create(
            model="text-davinci-003",
            prompt="请根据我给你的指令生成" + str(prompt_num) + "条相似的指令，指令如下:" + prompt,
            temperature=1,
            max_tokens=1000,
            top_p=1,
            frequency_penalty=0.0,
            presence_penalty=0.0,
            stop=None
        )
        # 拆分指令
        res_prompt = split_prompt(response, task)
        # 回调
        gpt_callback(res_prompt, task)
        logger.info('open_ai_结束')
    except Exception as e:
        logger.info(f"错误 {e}")


if __name__ == '__main__':
    key = 'laravel_database_gpt_text_prompt'
    logger.info('开始执行')
    while True:
        tasks = list()
        # 获取任务数量
        task_num = RedisCon.get_task_num(key)
        print(task_num, '个任务')
        nums = 10
        if task_num and task_num > 0:
            if task_num < nums:
                process = task_num
            else:
                # 最多同时10个任务
                process = nums
            loop = asyncio.get_event_loop()
            for i in range(process):
                task = RedisCon.get_task(key)
                if task:
                    # 存在任务--转对象
                    task = json.loads(task)
                    logger.info(task)
                    tasks.append(get_gpt_message(task))
            loop.run_until_complete((asyncio.wait(tasks, timeout=30)))
        else:
            logger.info('等待3秒')
            time.sleep(3)
            continue
    # test_task = {'type': 2,
    #              'prompt': 'a very cute boys, 16 years old, fly a kite,in beautifulstreet,full body，3d art, c4d, octane render, ray tracting,clay material,popmart blind box，Pixar trend, animationlighting, depth of field, ultra detailed',
    #              'prompt_img': None, 'user_id': 1, 'prompt_num': 5,
    #              'callback': 'http://test.phpgpt.com/api/users/gpt/callback', 'task_id': 1}
    # get_gpt_message(test_task)
