import requests
from PIL import Image
from io import BytesIO
import base64
import time
import os

# 新增函数：从api.txt获取服务器地址
def get_api_base_url():
    try:
        with open("api.txt", "r") as f:
            # 读取第一行有效内容（忽略空行和空白）
            lines = [line.strip() for line in f.readlines() if line.strip()]
            if lines:
                return lines[0].rstrip('/')  # 移除末尾的斜杠
    except FileNotFoundError:
        print("api.txt文件未找到，使用默认地址")
    except Exception as e:
        print(f"读取api.txt失败: {e}，使用默认地址")
    
    # 默认地址
    return "http://localhost:3000"

# 获取当前时间戳
def get_current_timestamp():
    return int(time.time())

# 获取unikey
def get_unikey():
    base_url = get_api_base_url()
    url = f"{base_url}/login/qr/key?time={get_current_timestamp()}"
    response = requests.get(url)
    data = response.json()
    if data['code'] == 200:
        return data['data']['unikey']
    return None

# 创建二维码
def create_qr(unikey):
    base_url = get_api_base_url()
    url = f"{base_url}/login/qr/create?key={unikey}&qrimg=1&time={get_current_timestamp()}"
    response = requests.get(url)
    data = response.json()
    if data['code'] == 200:
        qrimg_base64 = data['data']['qrimg']
        return qrimg_base64
    return None

# 显示二维码图像
def display_qr_image(qrimg_base64):
    img_data = base64.b64decode(qrimg_base64.split(",")[1])
    img = Image.open(BytesIO(img_data))
    img.show()

# 监控扫码状态
def check_scan_status(unikey):
    base_url = get_api_base_url()
    url = f"{base_url}/login/qr/check?key={unikey}&time={get_current_timestamp()}"
    response = requests.get(url)
    return response.json()

# 登录函数（保持原有实现不变）
def login():
    unikey = get_unikey()
    if not unikey:
        print("无法获取unikey")
        return None

    qrimg_base64 = create_qr(unikey)
    if not qrimg_base64:
        print("无法创建二维码")
        return None
    
    display_qr_image(qrimg_base64)

    print("请扫描二维码，等待扫码成功...")

    while True:
        status = check_scan_status(unikey)
        
        if status['code'] == 803:
            print("授权登录成功")
            cookie = status['cookie']
            with open("cookies.txt", "w") as f:
                f.write(cookie)
            return cookie
        elif status['code'] == 801:
            print("等待扫码，继续请求中...")
        elif status['code'] == 802:
            print("授权中，继续请求中...")
        else:
            print(f"扫码失败: {status['message']}")
            return None
        
        time.sleep(2)

# 使用示例
if __name__ == "__main__":
    login()