from tqdm import tqdm

from django.core.management.base import BaseCommand
from django.db import transaction

from lms.courses.models import CourseStudent
from lms.enrollments.models import EnrolledUser


class DryRunException(Exception):
    pass


class Command(BaseCommand):
    def delete_rejected_users(self, options):
        enrolled_user_course_student_ids = []
        course_students_to_delete_ids = []
        try:
            with transaction.atomic():
                hrdb_only = options.get('hrdb_only', False)
                rejected_users = EnrolledUser.objects.filter(status=EnrolledUser.StatusChoices.REJECTED)
                if hrdb_only:
                    rejected_users = rejected_users.filter(hrdb_id__isnull=False)
                for rejected_user in tqdm(rejected_users, total=rejected_users.count()):
                    cs = CourseStudent.objects.filter(
                        course=rejected_user.course,
                        group=rejected_user.group,
                        user=rejected_user.user,
                    ).first()
                    if cs is not None:
                        course_students_to_delete_ids.append(cs.id)
                        enrolled_user_course_student_ids.append((rejected_user.id, cs.id))

                CourseStudent.objects.filter(id__in=course_students_to_delete_ids).delete()

                if options.get('dry_run', False):
                    raise DryRunException

        except DryRunException:
            self.stdout.write('Rollback')
        else:
            self.stdout.write('Commit')

        self.stdout.write(
            f"""
                Students to delete: {len(enrolled_user_course_student_ids)}
                Deleted students (enrolled_user.id, course_student.id): {enrolled_user_course_student_ids}
            """
        )

    def add_arguments(self, parser):
        parser.add_argument(
            '--dry-run', action='store_true', dest='dry_run',
            help='Run without commit to db',
        )
        parser.add_argument(
            '--hrdb-only', action='store_true', dest='hrdb_only',
            help='Only for enrolled users from HRDB',
        )

    def handle(self, *args, **options):

        self.delete_rejected_users(options)

        self.stdout.write("Done\n")
