from django.core.management.base import BaseCommand
from django.db import connections
from establishment.management.commands.add_position import namedtuplefetchall
from main.models import AwardType
from location.models import Country
from tqdm import tqdm

class Command(BaseCommand):
    help = '''Add award types from old db to new db.
    Run after migrate country code!'''

    def award_types_sql(self):
        with connections['legacy'].cursor() as cursor:
            cursor.execute('''
                            SELECT 
                                DISTINCT
                                at.id, TRIM(at.title) AS name, 
                                s.country_code_2 AS country_code
                            FROM award_types as at
                            JOIN sites s on s.id = at.site_id
                            WHERE LENGTH(TRIM(at.title))>0               
                           ''')
            return namedtuplefetchall(cursor)

    def handle(self, *args, **kwargs):
        objects =[]
        for a in tqdm(self.award_types_sql(), desc='Add award types: '):
            country = Country.objects.filter(code=a.country_code).first()
            if country:
                type = AwardType(name=a.name, old_id=a.id)
                type.country = country
                objects.append(type)
        types = AwardType.objects.bulk_create(objects)
        self.stdout.write(self.style.WARNING(f'Created award types objects.'))