from django.core.management.base import BaseCommand
import logging
import pickle
import json
from user_management.models import GmailCredential
from oauth2client import client

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = '修复Gmail凭证数据格式问题'

    def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('开始修复Gmail凭证...'))
        credentials = GmailCredential.objects.all()
        
        fixed_count = 0
        error_count = 0
        
        for cred in credentials:
            try:
                if not cred.credentials:
                    self.stdout.write(self.style.WARNING(f'ID {cred.id} 的凭证为空，跳过'))
                    continue
                    
                # 检测当前凭证格式
                self.stdout.write(f'处理凭证 ID: {cred.id}, 邮箱: {cred.gmail_email}')
                
                # 1. 尝试作为JSON加载
                try:
                    if isinstance(cred.credentials, str):
                        # 验证是否为有效JSON
                        json.loads(cred.credentials)
                        self.stdout.write(self.style.SUCCESS(f'凭证 {cred.id} 已是有效JSON格式'))
                        fixed_count += 1
                        continue
                except Exception:
                    pass
                
                # 2. 尝试从二进制pickle加载并转换为JSON
                try:
                    # 处理可能的pickle格式
                    if isinstance(cred.credentials, str):
                        oauth_creds = pickle.loads(cred.credentials.encode('latin1'))
                    else:
                        oauth_creds = pickle.loads(cred.credentials)
                    
                    # 转换为JSON并保存
                    json_creds = oauth_creds.to_json()
                    cred.credentials = json_creds
                    cred.save()
                    
                    self.stdout.write(self.style.SUCCESS(f'凭证 {cred.id} 已从pickle转换为JSON格式'))
                    fixed_count += 1
                    continue
                except Exception as e:
                    self.stdout.write(self.style.ERROR(f'无法处理凭证 {cred.id}: {str(e)}'))
                    error_count += 1
                    
            except Exception as e:
                self.stdout.write(self.style.ERROR(f'处理凭证 {cred.id} 时出错: {str(e)}'))
                error_count += 1
                
        self.stdout.write(self.style.SUCCESS(f'处理完成! 成功: {fixed_count}, 失败: {error_count}'))