from django.db import migrations


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

    versions = []
    for offset, solution in enumerate(range(2, 22), start=180):
        coefficient = solution % 7 + 2
        constant = solution % 11 + 1
        total = coefficient * solution + constant
        question, _ = Question.objects.update_or_create(
            slug=f"a-auto-linear-{offset:04d}",
            defaults={
                "track": "advanced",
                "tags": ["口算"],
                "is_active": True,
            },
        )
        version, _ = QuestionVersion.objects.update_or_create(
            question=question,
            version=1,
            defaults={
                "prompt": f"{coefficient}x + {constant} = {total}，求 x",
                "answer": str(solution),
                "explanation": f"答案为 {solution}",
            },
        )
        versions.append(version)

    for contest in Contest.objects.filter(track="advanced").iterator():
        existing_ids = set(
            ContestQuestion.objects.filter(contest=contest).values_list(
                "question_version_id",
                flat=True,
            )
        )
        next_order = (
            ContestQuestion.objects.filter(contest=contest)
            .order_by("-order")
            .values_list("order", flat=True)
            .first()
            or 0
        )
        additions = []
        for version in versions:
            if version.id in existing_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)


class Migration(migrations.Migration):
    dependencies = [
        ("contest", "0007_remove_realtimematch_matchmaking_lookup_idx_and_more"),
    ]

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