from django.core.management.base import BaseCommand
from device_interaction.models import DeviceType, DeviceBatch, Device, UserDevice

class Command(BaseCommand):
    help = '清理测试用设备'

    def add_arguments(self, parser):
        parser.add_argument('--all', action='store_true', help='清理所有设备数据')
        parser.add_argument('--test-only', action='store_true', help='只清理名称中包含"测试"的设备数据')
        parser.add_argument('--force', action='store_true', help='强制删除，不提示确认')

    def handle(self, *args, **options):
        all_data = options['all']
        test_only = options['test_only']
        force = options['force']
        
        if not all_data and not test_only:
            self.stdout.write(self.style.ERROR('请指定清理范围: --all 或 --test-only'))
            return
        
        # 显示警告并请求确认
        if all_data:
            self.stdout.write(self.style.WARNING('警告: 即将删除所有设备数据!'))
        else:
            self.stdout.write(self.style.WARNING('警告: 即将删除所有测试设备数据!'))
        
        if not force:
            confirm = input('确认删除？(y/n): ')
            if confirm.lower() != 'y':
                self.stdout.write(self.style.SUCCESS('操作已取消'))
                return
        
        # 开始删除过程
        self.stdout.write('开始清理数据...')
        
        # 先删除用户设备关联
        user_device_query = UserDevice.objects.all()
        if test_only:
            test_devices = Device.objects.filter(device_type__name__contains='测试')
            user_device_query = user_device_query.filter(device__in=test_devices)
        
        user_device_count = user_device_query.count()
        user_device_query.delete()
        self.stdout.write(f'已删除 {user_device_count} 条用户设备关联')
        
        # 删除设备
        device_query = Device.objects.all()
        if test_only:
            device_query = device_query.filter(device_type__name__contains='测试')
        
        device_count = device_query.count()
        device_query.delete()
        self.stdout.write(f'已删除 {device_count} 台设备')
        
        # 删除批次
        batch_query = DeviceBatch.objects.all()
        if test_only:
            batch_query = batch_query.filter(device_type__name__contains='测试')
        
        batch_count = batch_query.count()
        batch_query.delete()
        self.stdout.write(f'已删除 {batch_count} 个批次')
        
        # 删除设备类型
        type_query = DeviceType.objects.all()
        if test_only:
            type_query = type_query.filter(name__contains='测试')
        
        type_count = type_query.count()
        type_query.delete()
        self.stdout.write(f'已删除 {type_count} 种设备类型')
        
        self.stdout.write(self.style.SUCCESS(
            f'清理完成! 总计删除:\n'
            f'- {user_device_count} 条用户设备关联\n'
            f'- {device_count} 台设备\n'
            f'- {batch_count} 个批次\n'
            f'- {type_count} 种设备类型'
        ))