from django.core.management.base import BaseCommand
from django.contrib.contenttypes.models import ContentType
from django.db import transaction
from hookmigration.models import Transient
from hookmigration.peewee_models import CustomersReferral
from django.contrib.auth.models import User
from progressbar import ProgressBar
from loyalty.models import Point
from referral.models import Referrer


class Command(BaseCommand):
    help = 'Migrate the users points'

    def handle(self, *args, **options):
        referral_records = CustomersReferral.select()

        total = referral_records.count()
        finished = 0
        progress = ProgressBar(max_value=total)

        user_content_type = ContentType.objects.get_for_model(User)
        referral_content_type = ContentType.objects.get_for_model(Referrer)

        missing = set()

        with transaction.atomic():
            for record in referral_records:
                new_user_transient = Transient.objects.filter(content_type=user_content_type).filter(
                    old_pk=record.user_id).first()
                if not new_user_transient:
                    missing.add(record.user_id)
                    continue
                user = User.objects.get(pk=new_user_transient.new_pk)

                # Add the point
                referral = Referrer.objects.create(
                    user=user,
                    code=record.code,
                )

                # Add the transient
                Transient.objects.create(
                    old_pk=record.id,
                    new_pk=referral.pk,
                    content_type=referral_content_type,
                    old_model=CustomersReferral.__class__.__name__
                )

                finished += 1
                progress.update(finished)
        print('All Done')