"""
补录 merchant_id 字段的 management command

用于为历史数据补充 merchant_id 字段值。
"""
from django.core.management.base import BaseCommand, CommandError
from django.db import transaction

from basic_info.models import Merchant
from printing.models import PlateOrder, PrintingOrder, PrintingJob


class Command(BaseCommand):
    help = '为 PlateOrder、PrintingOrder、PrintingJob 补录 merchant_id 字段'

    def add_arguments(self, parser):
        parser.add_argument(
            'merchant_id',
            type=int,
            help='要设置的商户ID'
        )
        parser.add_argument(
            '--dry-run',
            action='store_true',
            help='仅显示将要更新的记录数，不实际执行'
        )

    def handle(self, *args, **options):
        merchant_id = options['merchant_id']
        dry_run = options['dry_run']

        # 验证商户是否存在
        try:
            merchant = Merchant.objects.get(id=merchant_id)
        except Merchant.DoesNotExist:
            raise CommandError(f'商户 ID {merchant_id} 不存在')

        self.stdout.write(f'目标商户: {merchant.name} (ID: {merchant.id})')
        self.stdout.write('')

        # 统计需要更新的记录
        plate_orders_count = PlateOrder.objects.filter(merchant__isnull=True).count()
        printing_orders_count = PrintingOrder.objects.filter(merchant__isnull=True).count()
        printing_jobs_count = PrintingJob.objects.filter(merchant__isnull=True).count()

        self.stdout.write(f'待更新记录:')
        self.stdout.write(f'  - PlateOrder: {plate_orders_count} 条')
        self.stdout.write(f'  - PrintingOrder: {printing_orders_count} 条')
        self.stdout.write(f'  - PrintingJob: {printing_jobs_count} 条')
        self.stdout.write('')

        total = plate_orders_count + printing_orders_count + printing_jobs_count
        if total == 0:
            self.stdout.write(self.style.SUCCESS('没有需要更新的记录'))
            return

        if dry_run:
            self.stdout.write(self.style.WARNING('--dry-run 模式，未执行实际更新'))
            return

        # 执行更新
        with transaction.atomic():
            updated_plate_orders = PlateOrder.objects.filter(
                merchant__isnull=True
            ).update(merchant=merchant)
            
            updated_printing_orders = PrintingOrder.objects.filter(
                merchant__isnull=True
            ).update(merchant=merchant)
            
            updated_printing_jobs = PrintingJob.objects.filter(
                merchant__isnull=True
            ).update(merchant=merchant)

        self.stdout.write('')
        self.stdout.write(self.style.SUCCESS(f'更新完成:'))
        self.stdout.write(self.style.SUCCESS(f'  - PlateOrder: {updated_plate_orders} 条'))
        self.stdout.write(self.style.SUCCESS(f'  - PrintingOrder: {updated_printing_orders} 条'))
        self.stdout.write(self.style.SUCCESS(f'  - PrintingJob: {updated_printing_jobs} 条'))