import json
import time
import uuid
from multiprocessing import Process, Lock
import asyncio
from loguru import logger
from Helper import RedisCon
import requests
import os
from custom_oss_upload import AliyunOss


class DownLoad:
    def __init__(self, lock, task):
        self.key = ''
        self.lock = lock
        self.task = task
        self.header = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': ''
        }

    # 图片下载回调
    async def task_callback(self, status, result_img, message=''):
        url = self.task['callback']
        for i in range(5):
            try:
                params = {
                    'prompt_id': self.task['prompt_id'],
                }
                if status:
                    # 为真-成功
                    params['result_img'] = result_img
                else:
                    # 提交报错信息
                    params['message'] = message
                res = requests.post(url, data=params)
                logger.info(res)
                break
            except Exception as e:
                logger.error(e)

    async def download_img(self, url, cur_uuid):
        """
        下载图片
        :param url:
        :return:
        """
        for i in range(3):
            try:
                get_file = requests.get(url=url, headers=self.header, stream=True, allow_redirects=False, timeout=10)
                file_name = cur_uuid
                if not os.path.exists('./download'):
                    os.makedirs('./download')
                path = f'./download/{file_name}'
                with open(path, "wb") as Pypdf:
                    for chunk in get_file.iter_content(chunk_size=1024):  # 1024 bytes
                        if chunk:
                            Pypdf.write(chunk)
                return path
            except Exception as e:
                logger.error(e)

    async def upload_ali_oss(self, path, cur_uuid):
        # path--本地图片路径
        # name- 阿里云存储位置
        name = f"files/download/admin/{cur_uuid}"
        result_img = await AliyunOss().put_object_from_file(name, path)
        if result_img:
            # 回调
            await self.task_callback(True, result_img)

    # 开始
    async def on_start(self):
        logger.info('开始执行')
        cur_uuid = uuid.uuid4()
        cur_uuid = str(cur_uuid) + '.png'
        # 下载图片
        path = await self.download_img(self.task['result_img'], cur_uuid)
        if path:
            logger.info('图片下载成功')
            # 图片下载成功--准备上传阿里云
            await self.upload_ali_oss(path, cur_uuid)
        else:
            # 图片下载失败-直接回调-失败状态
            await self.task_callback(False, '', '转存discord图片失败')


def process_start(lock, task):
    loop = asyncio.get_event_loop()
    keyword = DownLoad(lock, task)
    # 生产
    loop.run_until_complete(keyword.on_start())


if __name__ == '__main__':
    # 获取任务并实例化
    key = 'laravel_database_discord_img_download'
    while True:
        try:
            tasks = list()
            task_num = RedisCon.get_task_num(key)
            logger.info(f'{task_num}个任务')
            nums = 5
            process = 0
            if task_num and task_num > 0:
                if task_num < nums:
                    process = task_num
                else:
                    process = nums
            lock = Lock()
            l = []
            for i in range(process):
                task = RedisCon.get_task(key)
                if task:
                    task = json.loads(task)
                    p = Process(target=process_start, args=(lock, task))
                    p.start()
                    l.append(p)  # 将进程加入列表中，这样 5 个子进程就会全部执行

            for p in l:
                p.join()  # 当5 个子进程全部执行完，才会执行主进程
        except Exception as e:
            logger.error(e)
            logger.info('等待3秒')
            time.sleep(3)
