from django.db import migrations


def expand_question_pools(apps, schema_editor):
    Contest = apps.get_model("contest", "Contest")
    ContestQuestion = apps.get_model("contest", "ContestQuestion")
    QuestionVersion = apps.get_model("contest", "QuestionVersion")

    for contest in Contest.objects.all().iterator():
        existing_version_ids = set(
            ContestQuestion.objects.filter(contest=contest).values_list(
                "question_version_id",
                flat=True,
            )
        )
        latest_versions = {}
        versions = QuestionVersion.objects.filter(
            question__track=contest.track,
            question__is_active=True,
        ).order_by("question_id", "-version")
        for version in versions.iterator():
            latest_versions.setdefault(version.question_id, version.id)

        next_order = (
            ContestQuestion.objects.filter(contest=contest)
            .order_by("-order")
            .values_list("order", flat=True)
            .first()
            or 0
        )
        additions = []
        for version_id in latest_versions.values():
            if version_id in existing_version_ids:
                continue
            next_order += 1
            additions.append(
                ContestQuestion(
                    contest_id=contest.id,
                    question_version_id=version_id,
                    order=next_order,
                    points=100,
                )
            )
        ContestQuestion.objects.bulk_create(additions, batch_size=500)


class Migration(migrations.Migration):
    dependencies = [
        ("contest", "0005_contestattempt_question_order_mathgameattempt_match_and_more"),
    ]

    operations = [
        migrations.RunPython(
            expand_question_pools,
            migrations.RunPython.noop,
        ),
    ]