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

from establishment.models import Establishment
from transfer.models import Establishments


class Command(BaseCommand):
    help = """
    Update establishment status field..
    Run command ./manage.py add_establishment_status
    """

    def get_status(self, old_status):
        if old_status:
            old_status = old_status.lower()
            if old_status == 'abandoned':
                return Establishment.ABANDONED
            elif old_status == 'closed':
                return Establishment.CLOSED
            elif old_status == 'published':
                return Establishment.PUBLISHED
            elif old_status == 'unpicked':
                return Establishment.UNPICKED
        return Establishment.WAITING

    def handle(self, *args, **kwargs):
        queryset = Establishment.objects.exclude(old_id__isnull=True)
        to_update = []
        non_existed_objects = 0

        for establishment in tqdm(queryset, desc='Iterate over existed establishments'):
            legacy_establishment_qs = Establishments.objects.filter(id=establishment.old_id)
            if legacy_establishment_qs.exists():
                legacy_establishment = legacy_establishment_qs.first()
                establishment.status = self.get_status(legacy_establishment.state)
                to_update.append(establishment)
            else:
                non_existed_objects += 1

        Establishment.objects.bulk_update(to_update, ['status', ])
        self.stdout.write(self.style.WARNING(f'Non existed {non_existed_objects} objects.'))
        self.stdout.write(self.style.SUCCESS(f'Updated {len(to_update)} objects.'))