from django.core.management.base import BaseCommand, CommandError
from products.models import *

class Command(BaseCommand):
    help = 'Update batch binding with tenant'

    def add_arguments(self, parser):
        parser.add_argument('--batch', '-b', type=int)
        parser.add_argument('--job', '-j', type=int)

    def handle(self, *args, **options):
        if options.get('job'):
            job = Job.objects.get(pk=options['job'])
        else:
            job = Job.objects.create(name='bind-batch.%d' % options.get('batch'))
        try:
            job.update('running', 0.0)
            b = CodeBatch.objects.get(pk=options['batch'])
            t = b.tenant
            q = SerialCode.objects.filter(batch=b)
            total = q.count()
            done = 0
            batchsize = 1000
            for objs in self.batch_iter(q, batchsize):
                if t:
                    seq_num = t.alloc_seq_nums(batchsize)
                    for x in objs:
                        x.tenant = t
                        x.seq_num = seq_num
                        seq_num += 1
                else:
                    for x in objs:
                        x.tenant = None
                        x.seq_num = None
                SerialCode.objects.bulk_update(objs, ['tenant', 'seq_num'])
                done += len(objs)
                perc = done * 100.0 / total
                job.update('running', perc)
                print(f"bound {done} codes, total {total}")
            job.update('done', 100.0)
        except Exception as e:
            job.update('error', message=str(e))

    def batch_iter(self, q, batchsize=1000):
        this_batch = []
        for x in q.iterator():
            this_batch.append(x)
            if len(this_batch) >= batchsize:
                yield this_batch
                this_batch = []
        if this_batch:
            yield this_batch