auto_cursor/services/email_manager.py

import asyncio
import email
from dataclasses import dataclass
from email.header import decode_header, make_header
from typing import Dict, List, Optional

import aiohttp
from loguru import logger
import aiomysql

from core.config import Config
from core.database import DatabaseManager
from core.exceptions import EmailError


@dataclass
class EmailAccount:
    id: int
    email: str
    password: str  # 这里实际上是 refresh_token
    client_id: str
    refresh_token: str
    in_use: bool = False
    cursor_password: Optional[str] = None
    cursor_cookie: Optional[str] = None
    sold: bool = False
    status: str = 'pending'  # 新增状态字段: pending, unavailable, success


class EmailManager:
    def __init__(self, config: Config, db_manager: DatabaseManager):
        self.config = config
        self.db = db_manager
        self.verification_subjects = [
            "Verify your email address",
            "Complete code challenge",
        ]
        # Redis相关配置
        self.use_redis = False
        if hasattr(self.db, 'redis') and self.db.redis:
            self.use_redis = True
            logger.info("Redis可用，将使用Redis进行邮箱状态管理")
        else:
            logger.warning("Redis不可用，将使用MySQL进行邮箱状态管理")
            
        # Redis键前缀
        self.redis_prefix = "emailmanager:"
        # 锁和黑名单的过期时间（秒）
        self.lock_timeout = 600  # 10分钟
        self.blacklist_timeout = 86400 * 30  # 30天
        # 黑名单初始化标记
        self.blacklist_initialized = False

    async def initialize(self):
        """初始化EmailManager，确保在使用前完成必要的设置"""
        if self.use_redis:
            await self._ensure_blacklist_initialized()
            self.blacklist_initialized = True

    async def _ensure_blacklist_initialized(self):
        """确保黑名单已经初始化"""
        if not self.use_redis:
            return False
            
        blacklist_key = f"{self.redis_prefix}blacklist:initialized"
        initialized = await self.db.redis.exists(blacklist_key)
        
        if not initialized:
            logger.info("初始化Redis邮箱黑名单...")
            # 查询所有已成功或不可用的邮箱
            query = """
                SELECT DISTINCT email 
                FROM email_accounts 
                WHERE status = 'success' OR status = 'unavailable'
            """
            results = await self.db.fetch_all(query)
            
            if results:
                # 批量添加到黑名单
                blacklist_key = f"{self.redis_prefix}blacklist:emails"
                emails = [row['email'] for row in results]
                if emails:
                    pipeline = self.db.redis.pipeline()
                    for email in emails:
                        pipeline.sadd(blacklist_key, email)
                    pipeline.expire(blacklist_key, self.blacklist_timeout)
                    await pipeline.execute()
                    
                    logger.info(f"已将 {len(emails)} 个邮箱添加到黑名单")
            
            # 标记为已初始化
            await self.db.redis.setex(f"{self.redis_prefix}blacklist:initialized", self.blacklist_timeout, "1")
            return True
        
        return True

    async def is_email_blacklisted(self, email: str) -> bool:
        """检查邮箱是否在黑名单中"""
        if not self.use_redis:
            # 回退到数据库查询
            query = """
                SELECT 1 FROM email_accounts 
                WHERE email = %s AND (status = 'success' OR status = 'unavailable')
                LIMIT 1
            """
            result = await self.db.fetch_one(query, (email,))
            return result is not None
            
        # 使用Redis SET存储黑名单
        blacklist_key = f"{self.redis_prefix}blacklist:emails"
        return await self.db.redis.sismember(blacklist_key, email)

    async def add_email_to_blacklist(self, email: str):
        """将邮箱添加到黑名单"""
        if not self.use_redis:
            return False
            
        blacklist_key = f"{self.redis_prefix}blacklist:emails"
        await self.db.redis.sadd(blacklist_key, email)
        await self.db.redis.expire(blacklist_key, self.blacklist_timeout)
        logger.debug(f"邮箱 {email} 已添加到黑名单")
        return True

    async def lock_account(self, account_id: int) -> bool:
        """锁定账号"""
        if not self.use_redis:
            # 如果不使用Redis，通过数据库更新来锁定
            try:
                query = """
                    UPDATE email_accounts 
                    SET in_use = 1, updated_at = CURRENT_TIMESTAMP
                    WHERE id = %s AND in_use = 0
                """
                affected = await self.db.execute(query, (account_id,))
                return affected > 0
            except Exception as e:
                logger.error(f"通过数据库锁定账号 {account_id} 失败: {e}")
                return False
                
        # 使用Redis实现分布式锁
        lock_key = f"{self.redis_prefix}lock:account:{account_id}"
        locked = await self.db.redis.setnx(lock_key, "1")
        if locked:
            # 锁定成功，设置过期时间
            await self.db.redis.expire(lock_key, self.lock_timeout)
            
            # 同时更新数据库状态
            try:
                update_query = """
                    UPDATE email_accounts 
                    SET in_use = 1, updated_at = CURRENT_TIMESTAMP
                    WHERE id = %s
                """
                await self.db.execute(update_query, (account_id,))
            except Exception as e:
                # 如果数据库更新失败，释放Redis锁
                logger.error(f"锁定账号 {account_id} 后更新数据库失败: {e}")
                await self.db.redis.delete(lock_key)
                return False
                
            logger.debug(f"账号 {account_id} 已锁定")
        return locked

    async def unlock_account(self, account_id: int) -> bool:
        """解锁账号"""
        # 无论是否使用Redis，都更新数据库
        try:
            update_query = """
                UPDATE email_accounts 
                SET in_use = 0, updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            """
            await self.db.execute(update_query, (account_id,))
        except Exception as e:
            logger.error(f"解锁账号 {account_id} 更新数据库失败: {e}")
            
        if not self.use_redis:
            return True
            
        # 删除Redis锁
        lock_key = f"{self.redis_prefix}lock:account:{account_id}"
        deleted = await self.db.redis.delete(lock_key)
        logger.debug(f"账号 {account_id} 锁已释放")
        return deleted > 0

    async def batch_get_accounts(self, num: int) -> List[EmailAccount]:
        """批量获取未使用的邮箱账号"""
        logger.info(f"尝试获取 {num} 个未使用的邮箱账号")
        
        # 如果使用Redis，确保黑名单已初始化
        if self.use_redis:
            await self._ensure_blacklist_initialized()
        
        # 1. 先从数据库中获取候选账号
        select_query = """
            SELECT id, email, password, client_id, refresh_token 
            FROM email_accounts 
            WHERE in_use = 0 AND sold = 0 AND status = 'pending'
            LIMIT %s
        """
        # 多获取一些候选账号，防止有些被排除
        candidate_accounts = await self.db.fetch_all(select_query, (num * 2,))
        
        if not candidate_accounts:
            logger.debug("没有找到符合条件的候选账号")
            return []
            
        logger.debug(f"找到 {len(candidate_accounts)} 个候选账号")
        
        # 2. 筛选并锁定账号
        result_accounts = []
        for account in candidate_accounts:
            # 检查邮箱是否在黑名单中
            if await self.is_email_blacklisted(account['email']):
                logger.debug(f"邮箱 {account['email']} 在黑名单中，跳过")
                continue
                
            # 尝试锁定账号
            if await self.lock_account(account['id']):
                # 添加到结果列表
                result_accounts.append(EmailAccount(
                    id=account['id'],
                    email=account['email'],
                    password=account['password'],
                    client_id=account['client_id'],
                    refresh_token=account['refresh_token'],
                    in_use=True
                ))
                
                # 如果已经获取足够的账号，退出循环
                if len(result_accounts) >= num:
                    break
            else:
                logger.debug(f"账号 {account['id']} 锁定失败，可能被其他进程使用")
        
        logger.info(f"实际获取到 {len(result_accounts)} 个可用账号")
        
        # 如果账号数量不足，尝试清理长时间锁定但未更新的账号
        if len(result_accounts) < num and len(result_accounts) < len(candidate_accounts):
            logger.warning("可用账号不足，尝试清理长时间锁定的账号")
            await self._cleanup_stuck_accounts()
            
        return result_accounts

    async def _cleanup_stuck_accounts(self):
        """清理长时间锁定但未更新的账号"""
        try:
            # 清理超过30分钟未更新且仍标记为in_use=1的账号
            cleanup_query = """
                UPDATE email_accounts 
                SET in_use = 0 
                WHERE in_use = 1 
                AND updated_at < DATE_SUB(NOW(), INTERVAL 30 MINUTE)
            """
            affected = await self.db.execute(cleanup_query)
            if affected > 0:
                logger.info(f"已清理 {affected} 个长时间锁定的账号")
                
                # 如果使用Redis，同时清理对应的锁
                if self.use_redis:
                    # 清理可能存在的Redis锁，但这需要知道具体的account_id
                    # 这里简化处理，依赖锁的自动过期机制
                    pass
        except Exception as e:
            logger.error(f"清理账号时出错: {e}")

    async def update_account_status(self, account_id: int, status: str):
        """更新账号状态"""
        try:
            # 获取账号邮箱信息（用于黑名单）
            account_query = "SELECT email, status as current_status FROM email_accounts WHERE id = %s"
            account_info = await self.db.fetch_one(account_query, (account_id,))
            
            if not account_info:
                logger.error(f"账号 {account_id} 不存在")
                return
                
            # 检查状态变更是否合理
            current_status = account_info.get('current_status')
            if current_status == 'success' and status != 'success':
                logger.warning(f"警告: 尝试将成功账号 {account_id} 状态改为 {status}，这可能是不正确的操作")
                # 如果已经是success状态，不允许降级为其他状态
                # return
            
            # 更新数据库状态
            query = '''
                UPDATE email_accounts 
                SET 
                    status = %s,
                    in_use = 0,
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            '''
            await self.db.execute(query, (status, account_id))
            logger.info(f"账号 {account_id} 状态已更新为 {status}")
            
            # 如果是success或unavailable状态，添加到黑名单
            if account_info and (status == 'success' or status == 'unavailable'):
                email = account_info['email']
                logger.debug(f"将邮箱 {email} 添加到黑名单 (状态: {status})")
                await self.add_email_to_blacklist(email)
            
            # 解锁账号
            await self.unlock_account(account_id)
            
            # 清除数据库缓存
            if self.db.redis:
                await self.db.clear_cache("db:*email_accounts*")
                
        except Exception as e:
            logger.error(f"更新账号 {account_id} 状态为 {status} 时出错: {e}")
            # 确保无论如何都解锁账号
            try:
                await self.unlock_account(account_id)
            except Exception as unlock_error:
                logger.error(f"尝试解锁账号 {account_id} 时出错: {unlock_error}")
            raise

    async def update_account(self, account_id: int, cursor_password: str, cursor_cookie: str, cursor_token: str):
        """更新账号信息（注册成功）"""
        try:
            # 获取账号邮箱信息（用于黑名单）
            account_query = "SELECT email FROM email_accounts WHERE id = %s"
            account_info = await self.db.fetch_one(account_query, (account_id,))
            
            # 更新数据库
            query = '''
                UPDATE email_accounts 
                SET 
                    cursor_password = %s, 
                    cursor_cookie = %s, 
                    cursor_token = %s,
                    in_use = 0, 
                    sold = 1,
                    status = 'success',
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            '''
            await self.db.execute(query, (cursor_password, cursor_cookie, cursor_token, account_id))
            logger.info(f"账号 {account_id} 更新为注册成功")
            
            # 添加到黑名单
            if account_info:
                email = account_info['email']
                logger.debug(f"将邮箱 {email} 添加到黑名单 (注册成功)")
                await self.add_email_to_blacklist(email)
            
            # 解锁账号
            await self.unlock_account(account_id)
            
            # 清除数据库缓存
            if self.db.redis:
                await self.db.clear_cache("db:*email_accounts*")
                
        except Exception as e:
            logger.error(f"更新账号 {account_id} 信息时出错: {e}")
            # 确保无论如何都解锁账号
            try:
                await self.unlock_account(account_id)
            except:
                pass
            raise

    async def release_account(self, account_id: int):
        """释放账号"""
        try:
            await self.unlock_account(account_id)
            logger.info(f"账号 {account_id} 已释放")
            
            # 清除数据库缓存
            if self.db.redis:
                await self.db.clear_cache("db:*email_accounts*")
                
        except Exception as e:
            logger.error(f"释放账号 {account_id} 时出错: {e}")
            raise

    async def count_pending_accounts(self) -> int:
        """统计可用的pending状态账号数量"""
        if self.use_redis:
            await self._ensure_blacklist_initialized()
            
        query = """
            SELECT COUNT(*) 
            FROM email_accounts 
            WHERE status = 'pending' AND in_use = 0 AND sold = 0
        """
        
        # 注：这里不使用黑名单过滤，因为数据量可能很大，
        # 但实际获取账号时会应用黑名单过滤
        
        result = await self.db.fetch_one(query)
        if result:
            return result.get("COUNT(*)", 0)
        return 0

    async def _get_access_token(self, client_id: str, refresh_token: str) -> str:
        """获取微软 access token"""
        logger.debug(f"开始获取 access token - client_id: {client_id}")

        url = 'https://login.microsoftonline.com/common/oauth2/v2.0/token'
        data = {
            'client_id': client_id,
            'grant_type': 'refresh_token',
            'refresh_token': refresh_token,
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(url, data=data) as response:
                result = await response.json()

                if 'error' in result:
                    error = result.get('error')
                    logger.error(f"获取 access token 失败: {error}")
                    raise EmailError(f"Failed to get access token: {error}")

                access_token = result['access_token']
                logger.debug("成功获取 access token")
                return access_token

    async def get_verification_code(self, email: str, refresh_token: str, client_id: str) -> str:
        """获取验证码"""
        logger.info(f"开始获取邮箱验证码 - {email}")
        try:
            # 1. 获取 access token
            access_token = await self._get_access_token(client_id, refresh_token)
            logger.debug(f"[{email}] 获取 access token 成功")

            # 2. 构建认证字符串
            auth_string = f"user={email}\1auth=Bearer {access_token}\1\1"
            logger.debug(f"[{email}] 认证字符串构建完成")

            # 3. 连接邮箱
            import imaplib
            mail = imaplib.IMAP4_SSL('outlook.live.com')
            mail.authenticate('XOAUTH2', lambda x: auth_string)
            mail.select('inbox')
            logger.debug(f"[{email}] 邮箱连接成功")

            # 4. 等待并获取验证码邮件
            for i in range(15):
                logger.debug(f"[{email}] 第 {i + 1} 次尝试获取验证码")

                # 搜索来自 no-reply@cursor.sh 的最新邮件
                result, data = mail.search(None, '(FROM "no-reply@cursor.sh")')
                if result != "OK" or not data[0]:
                    logger.debug(f"[{email}] 未找到来自 cursor 的邮件，等待1秒后重试")
                    await asyncio.sleep(1)
                    continue

                mail_ids = data[0].split()
                if not mail_ids:
                    logger.debug(f"[{email}] 邮件ID列表为空，等待1秒后重试")
                    await asyncio.sleep(1)
                    continue

                # 获取最新的3封邮件
                last_mail_ids = sorted(mail_ids, reverse=True)[:3]

                for mail_id in last_mail_ids:
                    result, msg_data = mail.fetch(mail_id, "(RFC822)")
                    if result != 'OK':
                        logger.warning(f"[{email}] 获取邮件内容失败: {result}")
                        continue

                    # 确保 msg_data 不为空且格式正确
                    if not msg_data or not msg_data[0] or len(msg_data[0]) < 2:
                        logger.warning(f"[{email}] 邮件数据格式不正确")
                        continue

                    # 正确导入 email 模块
                    from email import message_from_bytes
                    email_message = message_from_bytes(msg_data[0][1])

                    # 检查发件人
                    from_addr = str(make_header(decode_header(email_message['From'])))
                    if 'no-reply@cursor.sh' not in from_addr:
                        logger.debug(f"[{email}] 跳过非 Cursor 邮件，发件人: {from_addr}")
                        continue

                    # 检查主题
                    subject = str(make_header(decode_header(email_message['SUBJECT'])))
                    if not any(verify_subject in subject for verify_subject in self.verification_subjects):
                        logger.debug(f"[{email}] 跳过非验证码邮件，主题: {subject}")
                        continue

                    code = self._extract_code_from_email(email_message)
                    if code:
                        logger.debug(f"[{email}] 成功获取验证码: {code}")
                        mail.close()
                        mail.logout()
                        return code

                await asyncio.sleep(1)

            logger.error(f"[{email}] 验证码邮件未收到")
            raise EmailError("Verification code not received")

        except Exception as e:
            logger.error(f"[{email}] 获取验证码失败: {str(e)}")
            raise EmailError(f"Failed to get verification code: {str(e)}")

    def _extract_code_from_email(self, email_message) -> Optional[str]:
        """从邮件内容中提取验证码"""
        try:
            # 获取邮件内容
            if email_message.is_multipart():
                for part in email_message.walk():
                    if part.get_content_type() == "text/html":
                        body = part.get_payload(decode=True).decode('utf-8', errors='ignore')
                        break
            else:
                body = email_message.get_payload(decode=True).decode('utf-8', errors='ignore')

            # 提取6位数字验证码
            import re

            # 在HTML中查找包含6位数字的div
            match = re.search(r'<div[^>]*>(\d{6})</div>', body)
            if match:
                code = match.group(1)
                logger.debug(f"从HTML中提取到验证码: {code}")
                return code

            # 备用方案：搜索任何6位数字
            match = re.search(r'\b\d{6}\b', body)
            if match:
                code = match.group(0)
                logger.debug(f"从文本中提取到验证码: {code}")
                return code

            logger.warning(f"未能从邮件中提取到验证码")
            logger.debug(f"邮件内容预览: " + body[:200])
            return None

        except Exception as e:
            logger.error(f"提取验证码失败: {str(e)}")
            return None