from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support import expected_conditions as EC
import pyperclip
import time
from loguru import logger
from bit_api import *
import asyncio
from Helper import RedisCon
import json
from multiprocessing import Process
import random

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


class ToWork:
    def __init__(self):
        self.init_prompt = '/imagine'

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

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

    # 粘贴
    def driver_paste(self, dricer, value):
        ActionChains(dricer).send_keys(Keys.CONTROL, value).perform()

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

        return False


    # 等待元素出现--自定义
    def sleepElementShow(self, dricer, type, value):
        try:
            element = dricer.find_element(type, value)
            if element:
                return element
            return False
        except Exception as e:
            return False

    # 找元素，只找一次，不等待
    def FindElement(self, custom_el, find_type, value):
        try:
            el = custom_el.find_element(find_type, value)
            return el
        except Exception as e:
            return False

    # 获取元素列表
    def sleepElementListShow(self, dricer, type, value):
        try:
            element = dricer.find_elements(type, value)
            if element:
                return element
            return False
        except Exception as e:
            return False

    # 元素点击
    def ElementClick(self, dricer, type, value, name=''):
        try:
            task = dricer.find_element(type, value)
            if task:
                # 随机等待
                self.random_sleep(1, 2)
                # js 点击
                dricer.execute_script("arguments[0].click()", task)
                #  task.click()
                logger.info(f'{name}-元素存在，开始点击')
                return True
            logger.info(f'{name}-元素不存在')
            return False
        except Exception as e:
            logger.info(f'{name}-元素不存在')
            logger.error(e)
            return False

    # 循环输入指令
    def enter_prompt(self, dricer, tasks):
        if True:
            # 获取开始时间
            start_time = get_current_time()
            while True:
                # 再次获取当前时间
                current_time = get_current_time()
                # 最多等待一分钟
                if current_time - start_time >= 60:
                    # end
                    break
                # 开始判断
                element = '//ol[@data-list-id="chat-messages"]/li'
                # 消息列表
                message_list = self.CheckElement(dricer, 'xpath', element, 20, 'list')
                if message_list:
                    logger.info('消息列表存在')
                    # 找到所有子元素
                    for e in message_list:
                        logger.info(e.get_attribute('id'))
                        # 找到文本并匹配
                        path = '//div[contains(@data-list-item-id, "chat-messages___chat-messages")]'
                        label_element = self.FindElement(e, By.XPATH, path)
                        if label_element:
                            logger.info(label_element.get_attribute('data-list-item-id'))
                        else:
                            logger.info('未找到文本')
                logger.info('等待300秒')
                time.sleep(300)
        return
        # task是一个列表
        for task in tasks:
            # 输入初始指令
            logger.info('输入初始指令')
            self.driver_paste(dricer, self.init_prompt)
            # 选择指令
            prompt_el = '//*[@id="autocomplete-0"]'
            # 等待选择指令元素出现
            element = self.CheckElement(dricer, prompt_el, 5)
            if element:
                status = self.ElementClick(dricer, 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
                    self.driver_paste(dricer, change_prompt)
                    # 回车
                    ActionChains(dricer).send_keys(Keys.ENTER).perform()
                    time.sleep(2)
        # 输入完成,循环获取生成成功后的链接
        # 获取开始时间
        start_time = get_current_time()
        while True:
            # 再次获取当前时间
            current_time = get_current_time()
            # 最多等待一分钟
            if current_time - start_time >= 60:
                # end
                break
            # 开始判断
            message_list = '//*[@data-list-id="chat-messages"]'






    # 获取任务列表
    def get_current_tasks(self):


        # 测试
        # 当前浏览器需要执行的任务
        cur_tasks = []
        for j in range(2):
            # 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',
            }
            if task:
                # 将任务添加到列表
                cur_tasks.append(task)
        if len(cur_tasks):
            return cur_tasks
        return []

    def open_browser(self, browser_item):
        try:
            logger.info(f'浏览器id-{browser_item["num"]},开始执行')
            # /browser/open 接口会返回 selenium使用的http地址，以及webdriver的path，直接使用即可
            res = openBrowser(browser_item)
            driverPath = res['data']['driver']
            debuggerAddress = res['data']['http']
            chrome_options = Options()
            chrome_options.add_experimental_option("debuggerAddress", debuggerAddress)
            dricer = webdriver.Chrome(service=ChromeService(driverPath), options=chrome_options)
            logger.info('打开窗口')
            dricer.get('https://discord.com/channels/@me')
            # dricer.get('https://discord.com/login')
            return dricer
        except Exception as e:
            logger.error(e)
            return False

    # 初始化任务
    def on_start(self, browser_item):
        while True:
            # 先获取任务
            tasks = self.get_current_tasks()
            if len(tasks):
                # 打开窗口
                dricer = self.open_browser(browser_item)
                if not dricer:
                    # 浏览器打开失败
                    logger.error(f'{browser_item["num"]},窗口打开失败,检查代理ip')
                else:
                    # 开始执行任务
                    status = self.run_task(browser_item, tasks, dricer)
                    if status == False:
                        logger.info('请检查自动化')

            else:
                time.sleep(3)
                logger.info('不存在任务,等待3秒')

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


if __name__ == '__main__':
    key = 'laravel_database_midjourney_prompt'
    bro = ToWork()
    # while False:
    #     tasks = list()
    #     # 获取任务数量
    #     task_num = RedisCon.get_task_num(key)
    #     print(task_num, '个任务')
    #     nums = 3
    #     if task_num and task_num > 0:
    #         if task_num < nums:
    #             process = task_num
    #         else:
    #             # 最多同时3个任务
    #             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(bro.OpenBros(task))
    #         loop.run_until_complete((asyncio.wait(tasks, timeout=30)))
    #     else:
    #         time.sleep(3)
    #         continue


    # 测试-两个任务
    browser_list = [
        {
            # 7
            'num': 7,
            'id': '1115ce586624431a8b7bcba8b1e82035'
        },
        # {
        #     # 9
        #     'num': 9,
        #     'id': 'db474b6936a24a7a949377a6d90b4dcb'
        # }
    ]

    #
    l = []
    for item in browser_list:
        p = Process(target=bro.on_start, args=(item,))
        p.start()
        l.append(p)  # 将进程加入列表中，这样 5 个子进程就会全部执行

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



