import requests
import selenium
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import json
import base64

if __name__ == "__main__":
    url = "https://creator.xiaohongshu.com/login"
    
    driver = webdriver.Safari()
    driver.maximize_window()
    driver.get(url)
    time.sleep(5)
    WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.XPATH, "//img[@class='css-wemwzq']")))
    driver.find_element(By.XPATH, "//img[@class='css-wemwzq']").click()
    # time.sleep(10)

    # save img
    ## 等待跳转
    WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, "css-1lhmg90")))
    # 获取图片
    time.sleep(3)
    img_element = driver.find_element(By.CLASS_NAME, "css-1lhmg90")
    img_src = img_element.get_attribute("src")
    print(f"img_src: {img_src}")
    img_data = base64.b64decode(img_src.split(",")[1])
    print(img_data)
    with open("qrcode.png", "wb") as f:
        f.write(img_data)
            
    # 打印提示信息，请用户扫描二维码
    print("二维码已保存为 qrcode.png，请使用小红书 App 扫描登录")
    print("等待登录中...")
    
    # 保存当前 URL 用于检测页面跳转
    current_url = driver.current_url
    
    # 等待页面跳转（登录成功）
    # 方法1：检测 URL 变化
    try:
        # 使用一个较长的超时时间，给用户足够时间扫码
        WebDriverWait(driver, 120).until(lambda d: d.current_url != current_url)
        print("检测到页面跳转，登录成功！")
    except:
        # 如果 URL 没有变化，尝试方法2
        print("URL 没有变化，尝试检测登录状态元素...")
        
    # 方法2：检测登录后特定元素的出现
    # try:
    #     # 例如，检测用户头像或其他登录后才会出现的元素
    #     # 这里以用户头像为例，需要根据实际网页元素调整
    #     WebDriverWait(driver, 30).until(EC.presence_of_element_located((By.CLASS_NAME, "user-avatar")))
    #     print("检测到用户登录元素，登录成功！")
    # except:
    #     print("登录超时或失败，请重试")
    
    # 登录成功后，保存 cookies
    cookies = driver.get_cookies()
    with open("cookies_xhs.json", "w") as f:
        json.dump(cookies, f, indent=4)
    print(f"已保存 cookies 到 cookies_xhs.json 文件")
    
    # 关闭浏览器
    driver.quit()