#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Contact :   liuyuqi.gov@msn.cn
@Time    :   2025/09/30 22:31:02
@License :   Copyright © 2017-2022 liuyuqi. All Rights Reserved.
@Desc    :   alidriver_checkin script main entry
'''

import os
import argparse
from .crawl_lycheeshare import LycheeShare, DATA_DIR


def main():
    parser = argparse.ArgumentParser(description='LycheeShare 自动化工具')
    parser.add_argument('command', choices=['generate', 'register', 'checkin', 'token', 'credit'],
                        help='命令: generate(生成账户), register(批量注册), checkin(批量签到), token(批量获取token), credit(查询额度)')
    parser.add_argument('-n', '--number', type=int, default=1, help='生成账户数量 (用于generate命令)')

    args = parser.parse_args()

    if args.command == 'generate':
        lychee = LycheeShare('', '')
        lychee.generate_account(args.number)
        print(f"已生成 {args.number} 个账户到 {os.path.join(DATA_DIR, 'generate_account.txt')}")

    elif args.command == 'register':
        # 读取 generate_account.txt 批量注册
        generate_file = os.path.join(DATA_DIR, 'generate_account.txt')
        if not os.path.exists(generate_file):
            print(f"错误: 未找到 {generate_file} 文件，请先运行 generate 命令")
            return

        with open(generate_file, 'r') as f:
            lines = f.readlines()

        print(f"开始批量注册，共 {len(lines)} 个账户")
        success_count = 0

        for line in lines:
            parts = line.strip().split(',')
            if len(parts) < 3:
                print(f"跳过格式错误的行: {line.strip()}")
                continue

            username, email, password = parts[0], parts[1], parts[2]
            print(f"\n正在注册: {username} ({email})")

            lychee = LycheeShare(email, password)

            # 发送验证码（自动返回000000）
            success, result = lychee.get_captcha()
            if not success:
                print(f"发送验证码失败，跳过: {email}")
                continue

            # 自动使用验证码000000
            code = "000000"
            print(f"使用验证码: {code}")

            # 注册
            success, result = lychee.register(username, code)
            if success:
                # 注册成功，保存到 account.txt
                account_file = os.path.join(DATA_DIR, 'account.txt')
                with open(account_file, 'a') as af:
                    af.write(f"{username},{email},{password}\n")
                success_count += 1
                print(f"注册成功并保存到 {account_file}: {email}")

            # 添加延迟避免频繁请求
            import time
            import random
            time.sleep(10*60*random.random())  # 随机延迟1到10分钟

        print(f"\n批量注册完成！成功 {success_count}/{len(lines)} 个账户")

    elif args.command == 'checkin':
        # 读取 account.txt 批量签到
        account_file = os.path.join(DATA_DIR, 'account.txt')
        if not os.path.exists(account_file):
            print(f"错误: 未找到 {account_file} 文件，请先运行 register 命令")
            return

        with open(account_file, 'r') as f:
            lines = f.readlines()

        print(f"开始批量签到，共 {len(lines)} 个账户")
        success_count = 0

        for line in lines:
            parts = line.strip().split(',')
            if len(parts) < 3:
                print(f"跳过格式错误的行: {line.strip()}")
                continue

            username, email, password = parts[0], parts[1], parts[2]
            print(f"\n正在签到: {username} ({email})")

            lychee = LycheeShare(email, password)

            # 尝试加载cookie
            if lychee.load_cookies():
                print("使用已保存的登录状态")
            else:
                print("未找到登录状态，正在登录...")
                success, _ = lychee.login()
                if not success:
                    print(f"登录失败，跳过: {email}")
                    continue

            # 签到
            success, _ = lychee.checkin()
            if success:
                success_count += 1

            # 添加延迟避免频繁请求
            import time
            time.sleep(1)

        print(f"\n批量签到完成！成功 {success_count}/{len(lines)} 个账户")

    elif args.command == 'token':
        # 读取 account.txt 批量获取token
        account_file = os.path.join(DATA_DIR, 'account.txt')
        if not os.path.exists(account_file):
            print(f"错误: 未找到 {account_file} 文件，请先运行 register 命令")
            return

        with open(account_file, 'r') as f:
            lines = f.readlines()

        print(f"开始批量获取token，共 {len(lines)} 个账户")
        success_count = 0

        # 清空原有的 api_tokens.txt
        token_file = os.path.join(DATA_DIR, 'api_tokens.txt')
        if os.path.exists(token_file):
            os.remove(token_file)

        for line in lines:
            parts = line.strip().split(',')
            if len(parts) < 3:
                print(f"跳过格式错误的行: {line.strip()}")
                continue

            username, email, password = parts[0], parts[1], parts[2]
            print(f"\n正在获取token: {username} ({email})")

            lychee = LycheeShare(email, password)

            # 尝试加载cookie
            if lychee.load_cookies():
                print("使用已保存的登录状态")
            else:
                print("未找到登录状态，正在登录...")
                success, _ = lychee.login()
                if not success:
                    print(f"登录失败，跳过: {email}")
                    continue

            # 获取token
            success, _ = lychee.get_token()
            if success:
                success_count += 1

            # 添加延迟避免频繁请求
            import time
            time.sleep(1)

        print(f"\n批量获取token完成！成功 {success_count}/{len(lines)} 个账户")
        print(f"所有token已保存到 {os.path.join(DATA_DIR, 'api_tokens.txt')}")

    elif args.command == 'credit':
        # 读取 account.txt 批量查询额度
        account_file = os.path.join(DATA_DIR, 'account.txt')
        if not os.path.exists(account_file):
            print(f"错误: 未找到 {account_file} 文件，请先运行 register 命令")
            return

        with open(account_file, 'r') as f:
            lines = f.readlines()

        print(f"开始批量查询额度，共 {len(lines)} 个账户")
        success_count = 0

        # 清空原有的 credits.txt
        credit_file = os.path.join(DATA_DIR, 'credits.txt')
        if os.path.exists(credit_file):
            os.remove(credit_file)

        for line in lines:
            parts = line.strip().split(',')
            if len(parts) < 3:
                print(f"跳过格式错误的行: {line.strip()}")
                continue

            username, email, password = parts[0], parts[1], parts[2]
            print(f"\n正在查询额度: {username} ({email})")

            lychee = LycheeShare(email, password)

            # 尝试加载cookie
            if lychee.load_cookies():
                print("使用已保存的登录状态")
            else:
                print("未找到登录状态，正在登录...")
                success, _ = lychee.login()
                if not success:
                    print(f"登录失败，跳过: {email}")
                    continue

            # 查询额度
            success, _ = lychee.get_credit()
            if success:
                success_count += 1

            # 添加延迟避免频繁请求
            import time
            time.sleep(1)

        print(f"\n批量查询额度完成！成功 {success_count}/{len(lines)} 个账户")
        print(f"所有额度信息已保存到 {os.path.join(DATA_DIR, 'credits.txt')}")


if __name__ == '__main__':
    main()