import math

from django.core.management.base import BaseCommand
from celery import group
from establishment.models import Establishment
from establishment.tasks import update_establishment_image_urls


class Command(BaseCommand):
    help = """
    Updating image links for establishments.
    Run command ./manage.py update_establishment_image_urls
    """

    def add_arguments(self, parser):
        parser.add_argument(
            '--bucket_size',
            type=int,
            help='Size of one basket to update'
        )

    def handle(self, *args, **kwargs):
        bucket_size = kwargs.get('bucket_size', 128)

        objects = Establishment.objects.all()
        objects_size = objects.count()
        summary_tasks = math.ceil(objects_size / bucket_size)

        tasks = []

        for index in range(0, objects_size, bucket_size):
            bucket = objects[index: index + bucket_size]

            task = update_establishment_image_urls.s(
                (index + bucket_size) / bucket_size, summary_tasks,
                bucket.values_list('id', flat=True)
            )

            tasks.append(task)

        self.stdout.write(self.style.WARNING(f'Created all celery update tasks.\n'))

        job = group(*tasks)
        job.delay()

        self.stdout.write(self.style.SUCCESS(f'Done all celery update tasks.\n'))