import random

import undetected_chromedriver as uc
import time
import json


def get_cookies(browser, log_url="https://chatgpt.com/"):
    """
    获取cookies保存至本地
    """
    browser.get(log_url)
    input("回车以继续")
    # adkinsjoanna26@gmail.com
    # c1lO2NKEa2Hsl5

    dictCookies = browser.get_cookies()  # 获取list的cookies
    jsonCookies = json.dumps(dictCookies)  # 转换成字符串保存
    with open('damai_cookies.txt', 'w') as f:
        f.write(jsonCookies)
    print('cookies保存成功！')


def load_cookies(browser):
    """
    从本地读取cookies并刷新页面,成为已登录状态
    """
    with open('t.json', 'r', encoding='utf8') as f:
        listCookies = json.loads(f.read())

    # 往browser里添加cookies
    for cookie in listCookies:
        cookie_dict = {
            'domain': cookie.get('domain'),
            'name': cookie.get('name'),
            'value': cookie.get('value'),
            "expires": '',
            'path': '/',
            'httpOnly': False,
            'HostOnly': False,
            'Secure': False
        }
        try:
            browser.add_cookie(cookie_dict)
        except Exception as e:
            print("wrong_cookie: ", cookie_dict)
    browser.refresh()  # 刷新网页,cookies才成功


def get_presentation_are():
    # //div[@role="presentation"]//div[contains(@class, "text-sm") ]  # text-token-text-primary
    # //div[@role="presentation"]//div[contains(@class, "text-sm")]/div[contains(@class, "text-token-text-primary")]
    pass
    # //div[@role="presentation"]//p  # 响应的内容都放在了p里


def get_input_area(driver):
    input_area = driver.find_element(by="xpath", value="//textarea")
    return input_area


def get_input_button(driver):
    input_button = driver.find_element(by="xpath", value="//textarea/../../button")
    return input_button


def get_last_response(driver):
    # t = (driver.find_elements(by="xpath", value='//div[@role="presentation"]//p')[-1]).text
    t = (driver.find_elements(by="xpath",
                              value='//div[@role="presentation"]//div[contains(@class, "text-sm")]/div[contains(@class, "text-token-text-primary")]//div[contains(@class, "markdown")]')[
        -1]).text
    return t


def wait_for_complete(driver):
    time.sleep(3 + 3 * random.random())

    complete_flag = False
    while not complete_flag:
        try:
            last_bar = driver.find_elements(by='xpath', value="//div[contains(@class, 'mt-1 flex gap-3 empty:hidden juice:-ml-3')]")[-1]
            element_size = last_bar.size
            element_height = element_size['height']
            # element_width = element_size['width']
            if element_height < 10:
                time.sleep(5)
                print("sleep")
            else:
                complete_flag = True
        except Exception as e:
            time.sleep(5 + 5 * random.random())  # 第1次可能出现异常，可能因为第一次没有这个元素
            print('Exception')


def up_load_file(driver, file_name):
    file_input = driver.find_element(by="xpath", value='//input[@type="file"]')
    # 设置文件路径
    file_path = file_name
    file_input.send_keys(file_path)
    time.sleep(1 + random.random())

def sent_prompt(prompt, browser=None):
    input_area = get_input_area(browser)
    input_area.send_keys(prompt)
    time.sleep(1 + random.random())
    input_button = get_input_button(browser)
    input_button.click()
    wait_for_complete(driver)
    response_ = get_last_response(driver)
    return response_


def new_chat(driver):
    new_button = driver.find_element(by="xpath", value='//nav/div[1]/span[last()]/button')
    new_button.click()
    time.sleep(2 + random.random())


if __name__ == '__main__':
    driver = uc.Chrome()
    # get_cookies(driver)
    driver.get("https://chatgpt.com/")
    input("回车以登录")
    load_cookies(driver)
    input("回车以开始使用")
    file_name = r"C:\Users\zhu\Desktop\桌面整理\Zhou 等 - 2024 - TRAD Enhancing LLM Agents with Step-Wise Thought .pdf"
    up_load_file(driver, file_name)
    while True:
        prompt = input("输入prompt:")
        # prompt = "描写今天的天气很好，800字"
        response = sent_prompt(prompt, browser=driver)
        print(response)

    # # prompt = "描写今天的天气很好，800字"
    # response = sent_prompt(prompt, browser=driver)
    # print(response)
    driver.quit()