import requests as req
import pyperclip
import time
from loguru import logger
import asyncio
from Helper import RedisCon
import json
import random
import pyppeteer
from multiprocessing import Process, Lock
import uuid
import os

from Vmlogin import VmLogin


# 获取当前时间
def get_current_time():
    # 返回秒级时间戳
    return int(time.mktime(time.localtime(time.time())))


class ToWork:
    def __init__(self, browser_item, lock):
        self.browser_token = 'd8ee6138f2b7ae5e064061e5c29d611a'
        self.init_prompt = '/imagine'
        self.browser_item = browser_item
        self.browser_id = browser_item['id']
        self.lock = lock
        self.key = 'laravel_database_midjourney_prompt'
        self.header = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': ''
        }
        self.page = ''
        self.ws = ''

    # 随机等待
    def random_sleep(self, min, max):
        time.sleep(random.uniform(min, max))

    def click_time_rand(self, min_num, max_num):
        """
        点击随机时间
        :return:
        """
        return random.randint(min_num, max_num)

    # 粘贴板插入数据
    def driver_copy(self, driver, value):
        pyperclip.copy(value)

    # 粘贴--键盘输入
    async def driver_paste(self, page, value):
        for i in value:
            await page.keyboard.press(i)

    # 元素出现的循环次数
    async def CheckElement(self, page, path, num=10, check_type='only'):
        for i in range(num):
            if check_type == 'only':
                element = await self.sleepElementShow(page, path)
                if element:
                    # 随机等待
                    self.random_sleep(1, 2)
                    return element
            else:
                # 列表
                element = await self.sleepElementListShow(page, path)
                if element:
                    # 随机等待
                    self.random_sleep(1, 2)
                    return element
            time.sleep(1)

        return False

    # 等待元素出现--自定义
    async def sleepElementShow(self, page, path):
        try:
            element = (await page.xpath(path))[0]
            if element:
                return element
            return False
        except Exception as e:
            return False

    # 找元素，只找一次，不等待
    async def FindElement(self, custom_el, path, el_type=''):
        try:
            el = (await custom_el.xpath(path))
            # logger.info(len(el))
            # for i in el:
            #     title = await page.evaluate('(element) => element.textContent', i)
            #     logger.info(title)
            if el_type == 'last':
                return el[len(el) - 1]
            else:
                return el[0]
        except Exception as e:
            return False

    # 获取元素列表
    async def sleepElementListShow(self, page, path):
        try:
            element = (await page.xpath(path))
            if len(element):
                return element
            return False
        except Exception as e:
            logger.info(e)
            return False

    # 元素点击-xpath点击
    async def ElementClick(self, page, path, name=''):
        status = False
        for i in range(1, 10):
            try:
                ele = (await page.xpath(path))[0]
                await page.waitFor(self.click_time_rand(1, 2))
                await ele.click()
                await asyncio.sleep(self.click_time_rand(1, 2))
                logger.info(f'{name}-元素存在，开始点击')
                status = True
                break
            except Exception as e:
                logger.info(e)
                logger.info(f'{name}-元素不存在')
                status = False
                await asyncio.sleep(1)
        return status

    # 获取上传策略
    async def get_policy(self, url):
        for i in range(5):
            try:
                res = req.get(url)
                data = json.loads(res.text)['data']
                return data
            except Exception as e:
                logger.error('上传策略请求错误')
                logger.error(e)

    async def download_img(self, url, cur_uuid):
        """
        下载图片
        :param url:
        :return:
        """
        get_file = req.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

    async def image_to_byte(self, imagepath):
        '''
        图片转二进制
        '''
        with open(imagepath, "rb") as f:
            byte_data = f.read()
        return byte_data

    async def al_upload(self, policy, path, cur_uuid):
        # 阿里云上传
        url = 'https://' + policy['host']
        file = await self.image_to_byte(path)
        logger.info('图片已经转为二进制')
        params = {
            'key': policy['dir'] + cur_uuid,
            'policy': policy['policy'],
            'OSSAccessKeyId': policy['accessid'],
            'success_action_status': '200',
            'callback': policy['callback'],
            'signature': policy['signature'],
            'file': file,
        }
        header = {
            'Content-Type': 'multipart/form-data;charset=utf-8'
        }
        logger.info('阿里云上传')
        res = req.post(url, data=params, headers=header)
        logger.info(res)
        # 上传后的链接为
        upload_img = policy['domain'] + policy['dir'] + cur_uuid
        logger.info(upload_img)

    async def upload_img(self, task, href):
        try:
            cur_uuid = uuid.uuid4()
            cur_uuid = str(cur_uuid) + '.png'
            # 先获取上传策略
            policy = await self.get_policy(task['policy'])
            # 下载discord的图片
            # path = await self.download_img(href, cur_uuid)
            path = './download/0f24e1f6-4ba2-42d7-bdf4-852653f30fe4.png'
            logger.info(path)
            # 开始上传
            await self.al_upload(policy, path, cur_uuid)
        except Exception as e:
            logger.error(e)

    async def img_callback(self, task, href):
        # 先要上传图片
        # cur_href = await self.upload_img(task, href)
        url = task['callback']
        try:
            params = {
                'result_img': href,
                'user_id': task['user_id'],
                'task_id': task['task_id'],
                'prompt_id': task['prompt_id'],
            }
            res = req.post(url, data=params)
        except Exception as e:
            logger.error(e)

    # 检测图片生成进度
    async def check_img_load_process(self, page, task):
        logger.info('开始检测任务状态')
        # 先写单个任务的，获取最后一张图片
        while True:
            # 获取开始时间
            start_time = get_current_time()
            while True:
                # 再次获取当前时间
                current_time = get_current_time()
                # 最多等待一分钟
                if current_time - start_time >= 60:
                    # end
                    break
                cur_type = 1  # 获取最后一个
                if cur_type == 1:
                    logger.info('进入')
                    # 判断文字中是否存在进度条
                    text_path = '//div[contains(@id, "message-content-")]'
                    text_el = await self.FindElement(page, text_path, 'last')
                    if text_el:
                        title = await page.evaluate('(element) => element.textContent', text_el)
                        logger.info(title)
                        if '%' not in title and '(Waiting to start)' not in title and '正在发送命令' not in title:
                            logger.info('任务完成')
                            # 最后一个链接
                            img_path = '//*[contains(@id, "message-accessories-")]/div[1]/div/div/div/div/div/a'
                            img_el = await self.CheckElement(page, img_path, 20)
                            if img_el:
                                href = await page.evaluate('img_el => img_el.href', img_el)
                                if href:
                                    logger.info(href)
                                    # 回调
                                    await self.img_callback(task, href)

                            return
                        else:
                            logger.info('任务进行中')
                    else:
                        logger.info('文本元素不存在')
                else:
                    # 开始判断
                    element = '//ol[@data-list-id="chat-messages"]/li'
                    # 消息列表
                    message_list = await self.CheckElement(page, element, 20, 'list')
                    if message_list:
                        logger.info('消息列表存在')
                        # 找到所有子元素
                        for e in message_list:
                            # 找到文本并匹配
                            path = '//div[contains(@id, "message-content-")]/strong'
                            label_element = await self.FindElement(e, path)
                            if label_element:
                                title = await page.evaluate('(element) => element.textContent', label_element)
                                logger.info(title)
                            else:
                                logger.info('未找到文本')
                await asyncio.sleep(2)

    # 循环输入指令
    async def enter_prompt(self, page, tasks):
        # task是一个列表
        for task in tasks:
            # 开始判断
            element = '//ol[@data-list-id="chat-messages"]/li'
            # 消息列表
            message_list = await self.CheckElement(page, element, 20, 'list')
            if message_list:
                # 输入初始指令
                logger.info('输入初始指令')
                await self.driver_paste(page, self.init_prompt)
                # 选择指令
                prompt_el = '//*[@id="autocomplete-0"]'
                # 等待选择指令元素出现
                element = await self.CheckElement(page, prompt_el, 20)
                if element:
                    status = await self.ElementClick(page, prompt_el, '选择指令')
                    if status:
                        change_prompt = task['prompt']
                        # 图片链接
                        prompt_img = task['prompt_img']
                        if task['type'] == 1 and prompt_img and len(json.loads(prompt_img)):
                            prompt_img = json.loads(prompt_img)
                            st = ''
                            for item in prompt_img:
                                # 加入图片-链接
                                st += ' ' + item
                            change_prompt = st + ' ' + change_prompt
                        await self.driver_paste(page, change_prompt)
                        # 回车
                        await page.keyboard.press('Enter')
                        await asyncio.sleep(2)
                        # 输入完成,循环获取生成成功后的链接
                        await self.check_img_load_process(page, tasks[0])

    # 获取任务列表
    def get_current_tasks(self):
        if True:
            task = RedisCon.get_task(self.key)
            # 当前浏览器需要执行的任务
            cur_tasks = []
            if task:
                task = json.loads(task)
                task['prompt'] = task['prompt'].replace('，', ',')
                # 将任务添加到列表
                cur_tasks.append(task)
            if len(cur_tasks):
                return cur_tasks
            return []
        else:
            # 测试
            # 当前浏览器需要执行的任务
            cur_tasks = []
            for j in range(1):
                # task = RedisCon.get_task(key)
                task = {
                    '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 --q 5",
                    'user_id': 1,
                    'task_id': 1,
                    'prompt_img': [],
                    'result_img': "",
                    'type': 2,
                    'callback': 'http://test.phpgpt.com/api/users/prompt/callback',
                    'policy': 'http://test.phpgpt.com/api/users/config/policy2'
                }
                task['prompt'] = task['prompt'].replace('，', ',')
                if task:
                    # 将任务添加到列表
                    cur_tasks.append(task)
            if len(cur_tasks):
                return cur_tasks
            return []

    async def open_browser(self, browser_item):
        while True:
            try:
                logger.info(f'浏览器id-{browser_item["num"]},开始执行')
                vm = VmLogin(self.browser_token)
                proxy = vm.start_chrome(browser_item['id'])
                ws = ''
                if not proxy and self.ws:
                    logger.info('ws已存在')
                    ws = self.ws
                else:
                    url = f'http://127.0.0.1:{proxy}/json/version'
                    res = req.get(url).json()
                    ws = res['webSocketDebuggerUrl']
                    self.ws = res['webSocketDebuggerUrl']
                browser = await pyppeteer.connect({
                    "browserWSEndpoint": ws,
                    'args': [
                        '--window-size=1536,824',
                        '--start-maximized',
                        # '--blink-settings=imagesEnabled=false',
                        '--n-images',
                        '--disable-background-timer-throttling',
                        '--disable-backgrounding-occluded-windows',
                        '--disable-renderer-backgrounding',
                    ]
                })
                if self.page:
                    await self.page.goto('https://discord.com/channels/@me')
                else:
                    page = await browser.newPage()
                    self.page = page
                    await self.page.goto('https://discord.com/channels/@me')
                return self.page


            except Exception as e:
                logger.error(e)
                logger.info('浏览器打开失败,正在重试')

    # 初始化任务
    async def on_start(self):
        # 打开窗口
        page = await self.open_browser(self.browser_item)
        if not page:
            # 浏览器打开失败
            logger.error(f'{self.browser_item["num"]},窗口打开失败')
            return
        else:
            logger.info(f'{self.browser_item["num"]},打开窗口')
        while True:
            # 获取任务
            tasks = self.get_current_tasks()
            if len(tasks):
                # 开始执行任务
                logger.info('开始执行任务')
                status = await self.run_task(self.browser_item, tasks, page)
            else:
                time.sleep(3)
                logger.info('不存在任务,等待3秒')

    async def run_task(self, browser_item, tasks, page):
        try:
            logger.info('等待登录成功')
            # 判断页面是否加载完成
            page_load = '//*[@data-list-item-id="guildsnav___home"]'
            element = await self.CheckElement(page, page_load, 30)
            if not element:
                logger.error('登录失败')
                # 去登陆
                return False
            # 登录成功
            logger.info('登录成功')
            # 等待服务器元素出现
            server_el = '//div[contains(@aria-label, "测试服务器")]'
            element = await self.CheckElement(page, server_el, 20)
            if element:
                # 点击指定服务器
                status = await self.ElementClick(page, server_el, '测试服务器')
                if status:
                    # 等待聊天窗口出现
                    input_el = '//*[@role="textbox"]'
                    element = await self.CheckElement(page, input_el)
                    if element:
                        logger.info('聊天窗口存在')
                        # 输入指令
                        await self.enter_prompt(page, tasks)
                        return True
            return False
        except Exception as e:
            logger.error(e)
            return False


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

    # 测试
    # task = {
    #     'policy': 'http://test.phpgpt.com/api/users/config/policy2'
    # }
    # href = 'https://cdn.discordapp.com/attachments/1090583268467945475/1093734564188409946/suifeng_a_very_cute_boys_16_years_old_fly_a_kitein_beautifulstr_cd2d65d6-ea40-432b-9941-c796504b6082.png'
    # loop.run_until_complete(keyword.upload_img(task, href))


if __name__ == '__main__':
    browser_list = [
        {
            'num': 2,
            'id': 'F1670ACB-79E6-4071-A7EA-69ECDA448ED8'
        }
    ]
    lock = Lock()
    l = []
    for item in browser_list:
        p = Process(target=process_start, args=(item, lock))
        p.start()
        l.append(p)  # 将进程加入列表中，这样 5 个子进程就会全部执行

    for p in l:
        p.join()  # 当5 个子进程全部执行完，才会执行主进程
