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


class ImageSplit:
    def __init__(self, lock, task):
        self.header = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': ''
        }
        self.lock = lock
        self.task = task
        # 切割后的图片目录
        self.split_path = './split_assets'
        self.baidu_host = 'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=SEVtEVlQPyqf55mgGt5u5Ad2&client_secret=0CNGhjMrgXRPmDE014S2IXyHLKCuTcfK'
        # 百度接口的token
        self.baidu_token = ''

    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 mkdir(self, path):
        path = path.strip()
        path = path.rstrip("\\")

        isExists = os.path.exists(path)
        if not isExists:
            os.makedirs(path)
            print(path + ' 创建成功')
            return True
        else:
            print(path + ' 目录已存在')
            return False

    # 切割图片
    async def split_img(self, img_file):
        img = cv2.imread(img_file)

        # 这一段是改为灰色
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        img_h = img.shape[0]  # 高度
        img_w = img.shape[1]  # 宽度
        # h1_half = int(img_h / 2)
        # w1_half = int(img_w / 2)

        h1_half = img_h // 2
        w1_half = img_w // 2

        img_name = os.path.basename(img_file)
        for i in range(4):
            img1 = img[int(i / 2) * h1_half: h1_half * (int(i / 2) + 1),
                   int(i % 2) * w1_half: (int(i % 2) + 1) * w1_half]

            img1_path = os.path.join(self.split_path, f"{img_name[:-4]}_{i + 1}.png")
            cv2.imwrite(img1_path, img1)

    # 文件是否存在
    async def file_isexist(self, file_name, cut_id, path):
        read_file_name = f'{path}/{str(file_name)}_{cut_id}.png'
        exist = os.path.exists(read_file_name)
        if exist:
            return True
        return False

    # 获取百度的token
    async def get_api_token(self):
        try:
            for i in range(5):
                response = requests.get(self.baidu_host)
                if response:
                    access_token = response.json()['access_token']
                    self.baidu_token = access_token
                    return True
        except Exception as e:
            logger.error(e)
        return False

    async def Img_enhance(self, path):
        '''
        图片增强
        '''
        request_url = "https://aip.baidubce.com/rest/2.0/image-process/v1/image_definition_enhance"
        # 二进制方式打开图片文件
        f = open(path, 'rb')
        img = base64.b64encode(f.read())

        params = {"image": img}
        request_url = request_url + "?access_token=" + self.baidu_token
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        if response:
            imgdata = base64.b64decode(response.json()["image"])
            file = open(path, 'wb')
            file.write(imgdata)
            file.close()

    async def Img_enlarge(self, path):
        '''
        图片无损放大
        '''
        request_url = "https://aip.baidubce.com/rest/2.0/image-process/v1/image_quality_enhance"
        # 二进制方式打开图片文件
        f = open(path, 'rb')
        img = base64.b64encode(f.read())

        params = {"image": img}
        request_url = request_url + "?access_token=" + self.baidu_token
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        if response:
            imgdata = base64.b64decode(response.json()["image"])
            file = open(f'{path}', 'wb')
            file.write(imgdata)
            file.close()

    # 任务回调
    async def task_callback(self, result_img):
        url = self.task['callback']
        for i in range(5):
            try:
                res = requests.post(url, data={
                    'cut_img': result_img,
                    'prompt_id': self.task['prompt_id']
                })
                back = res.json()
                logger.info(back)
                if back['code'] == 0:
                    break
            except Exception as e:
                logger.error(e)

    # 开始
    async def on_start(self):
        download_url = self.task['result_img']
        cur_uuid = uuid.uuid4()
        new_uuid = str(cur_uuid) + '.png'
        path = await self.download_img(download_url, new_uuid)
        if path:
            # 下载成功
            # 创建目录
            await self.mkdir(self.split_path)
            # 开始切割
            await self.split_img(os.path.join('./download', new_uuid))
            logger.info('图片切割完成')
            # 判断文件是存在
            status = await self.file_isexist(cur_uuid, self.task['cut_id'], self.split_path)
            if status:
                logger.info('切割文件存在')
                # 获取百度接口的token
                baidu_status = await self.get_api_token()
                if not baidu_status:
                    logger.info('百度token获取失败')
                else:
                    # 文件路径
                    file_path = f'{self.split_path}/{str(cur_uuid)}_{self.task["cut_id"]}.png'
                    # 开始图像无损放大
                    logger.info('开始图像无损放大')
                    await self.Img_enlarge(file_path)
                    # 开始图像质量增强
                    await self.Img_enhance(file_path)
                    logger.info('图像增强结束')
                    # 上传阿里云
                    name = f"files/download/admin/{new_uuid}"
                    result_img = AliyunOss().put_object_from_file(name, file_path)
                    if result_img:
                        # 回调
                        await self.task_callback(result_img)
                        logger.info('本次任务结束')
            else:
                logger.info('切割文件不存在')


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


if __name__ == '__main__':
    # 获取任务并实例化
    key = 'laravel_database_discord_img_split'
    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 个子进程全部执行完，才会执行主进程
            time.sleep(3)
        except Exception as e:
            logger.error(e)
            logger.info('等待3秒')
            time.sleep(3)
