from django.core.management.base import BaseCommand
from tqdm import tqdm

from establishment.models import Establishment
from transfer.models import Establishments
from transfer.serializers.establishment import EstablishmentSerializer
from timetable.models import Timetable
from django.db import transaction


class Command(BaseCommand):
    help = 'Fix scheduler'

    @transaction.atomic
    def handle(self, *args, **kwargs):
        count = 0
        establishments = Establishment.objects.all()
        old_est_list = Establishments.objects.prefetch_related(
            'schedules_set',
        )
        # remove old records of Timetable
        Timetable.objects.all().delete()

        for est in tqdm(establishments, desc="Fix scheduler"):
            old_est = old_est_list.filter(id=est.old_id).first()

            if old_est and old_est.schedules_set.exists():
                old_schedule = old_est.schedules_set.first()
                timetable = old_schedule.timetable
                if timetable:
                    new_schedules = EstablishmentSerializer.get_schedules(timetable)
                    est.schedule.add(*new_schedules)
                    est.save()
                    count += 1

        self.stdout.write(self.style.WARNING(f'Update {count} objects.'))