from django.core.management.base import BaseCommand, CommandError
from user_management.models import GmailCredential, User
from user_management.gmail_integration import GmailIntegration
import logging
import pickle

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = '更新Gmail凭证中的邮箱信息'

    def add_arguments(self, parser):
        parser.add_argument('--email', type=str, help='指定用户邮箱')
        parser.add_argument('--gmail', type=str, help='指定要设置的Gmail邮箱')
        parser.add_argument('--all', action='store_true', help='更新所有凭证')

    def handle(self, *args, **options):
        email = options.get('email')
        gmail = options.get('gmail')
        update_all = options.get('all')
        
        if update_all:
            # 更新所有凭证
            credentials = GmailCredential.objects.filter(is_active=True)
            self.stdout.write(f"找到 {credentials.count()} 个活跃的Gmail凭证")
            
            for credential in credentials:
                self._update_credential(credential, gmail)
        elif email:
            # 更新指定用户的凭证
            try:
                user = User.objects.get(email=email)
                credentials = GmailCredential.objects.filter(user=user, is_active=True)
                
                if not credentials.exists():
                    raise CommandError(f"未找到用户 {email} 的Gmail凭证")
                
                for credential in credentials:
                    self._update_credential(credential, gmail)
            except User.DoesNotExist:
                raise CommandError(f"未找到用户 {email}")
        else:
            self.stdout.write("请提供--email参数或--all参数")
    
    def _update_credential(self, credential, gmail=None):
        """更新单个凭证"""
        user = credential.user
        self.stdout.write(f"正在更新用户 {user.email} 的Gmail凭证...")
        
        if gmail:
            # 如果指定了Gmail邮箱，直接使用
            credential.gmail_email = gmail
            credential.save()
            self.stdout.write(self.style.SUCCESS(f"已手动设置Gmail邮箱为: {gmail}"))
            return
        
        # 尝试使用API获取Gmail邮箱
        try:
            # 从凭证数据中恢复服务
            creds = pickle.loads(credential.credentials)
            
            if creds and not creds.invalid:
                # 创建Gmail集成实例
                integration = GmailIntegration(user=user)
                integration.credentials = creds
                
                # 尝试调用API获取用户资料
                profile = integration.gmail_service.users().getProfile(userId='me').execute()
                gmail_email = profile.get('emailAddress')
                
                if gmail_email:
                    credential.gmail_email = gmail_email
                    credential.save()
                    self.stdout.write(self.style.SUCCESS(f"已更新Gmail邮箱为: {gmail_email}"))
                else:
                    self.stdout.write(self.style.WARNING("无法从API获取Gmail邮箱"))
            else:
                self.stdout.write(self.style.ERROR("凭证无效，请重新授权"))
        except Exception as e:
            self.stdout.write(self.style.ERROR(f"更新失败: {str(e)}"))