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

from establishment.models import Establishment, EstablishmentSubType, EstablishmentType
from transfer.models import Metadata


class Command(BaseCommand):
    help = 'Add subtype for establishment artisan'

    def handle(self, *args, **options):
        artisans = Establishment.objects.artisans().filter(
            old_id__isnull=False,
        ).prefetch_related('tags')

        old_tags = Metadata.objects.filter(
            establishment__in=list(artisans.values_list('old_id', flat=True)),
            key='shop_category',
        )

        tags = []
        for tag in tqdm(old_tags):
            tags.append(tag.value)
        subtypes = set(tags)

        es_type, _ = EstablishmentType.objects.get_or_create(
            index_name='artisan',
            defaults={
                'index_name': 'artisan',
                'name': {'en-GB': 'artisan'},
            }
        )
        for artisan in tqdm(artisans):
            artisan_tags = artisan.tags.all()
            for t in artisan_tags:
                if t.value in subtypes:
                    tag = 'coffee_shop' if t.value == 'coffe_shop' else t.value
                    subtype, _ = EstablishmentSubType.objects.get_or_create(
                        index_name=tag,
                        defaults={
                            'index_name': tag,
                            'name': {'en-GB': ' '.join(tag.split('_')).capitalize()},
                            'establishment_type': es_type,
                        }
                    )
                    artisan.establishment_subtypes.add(subtype)
                    artisan.save()

        self.stdout.write(self.style.WARNING(f'Artisans subtype updated.'))
