# Generated by Django 5.0.9 on 2024-10-07 17:05

from django.db import migrations, transaction
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.db.models import F, Max, Min, OuterRef


def set_default_value_for_can_add_members_group(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    NamedUserGroup = apps.get_model("zerver", "NamedUserGroup")
    Realm = apps.get_model("zerver", "Realm")

    BATCH_SIZE = 1000
    max_id = NamedUserGroup.objects.filter(can_add_members_group=None).aggregate(Max("id"))[
        "id__max"
    ]

    if max_id is None:
        # Do nothing if there are no user groups on the server.
        return

    lower_bound = NamedUserGroup.objects.filter(can_add_members_group=None).aggregate(Min("id"))[
        "id__min"
    ]

    owners_system_group_ids = NamedUserGroup.objects.filter(name="role:owners").values_list(
        "id", flat=True
    )
    realm_ids_allowing_owners_to_manage_all_groups = Realm.objects.filter(
        can_manage_all_groups__in=owners_system_group_ids
    ).values_list("id", flat=True)

    admins_system_group_ids = NamedUserGroup.objects.filter(name="role:administrators").values_list(
        "id", flat=True
    )
    realm_ids_allowing_admins_to_manage_all_groups = Realm.objects.filter(
        can_manage_all_groups__in=admins_system_group_ids
    ).values_list("id", flat=True)

    moderators_system_group_ids = NamedUserGroup.objects.filter(name="role:moderators").values_list(
        "id", flat=True
    )
    realm_ids_allowing_moderators_to_manage_all_groups = Realm.objects.filter(
        can_manage_all_groups__in=moderators_system_group_ids
    ).values_list("id", flat=True)

    members_system_group_ids = NamedUserGroup.objects.filter(name="role:members").values_list(
        "id", flat=True
    )
    realm_ids_allowing_members_to_manage_all_groups = Realm.objects.filter(
        can_manage_all_groups__in=members_system_group_ids
    ).values_list("id", flat=True)

    while lower_bound <= max_id:
        upper_bound = lower_bound + BATCH_SIZE - 1
        print(f"Processing batch {lower_bound} to {upper_bound} for NamedUserGroup")

        with transaction.atomic():
            # Initialize to the corresponding system group, i.e. just
            # copy from can_manage_all_groups. Previously, either
            # owners/administrators/moderators or group members were
            # allowed to manage the group if they were part of
            # can_manage_all_groups. Non members were not allowed to
            # manage the group unless for the roles mentioned above.
            # That is why copying from can_manage_all_groups should
            # give the identical set of users permission.
            NamedUserGroup.objects.filter(
                id__range=(lower_bound, upper_bound),
                can_add_members_group=None,
                realm_id__in=realm_ids_allowing_owners_to_manage_all_groups,
            ).update(
                can_add_members_group=NamedUserGroup.objects.filter(
                    name="role:owners",
                    realm_for_sharding=OuterRef("realm_for_sharding"),
                    is_system_group=True,
                ).values("pk")
            )

            NamedUserGroup.objects.filter(
                id__range=(lower_bound, upper_bound),
                can_add_members_group=None,
                realm_id__in=realm_ids_allowing_admins_to_manage_all_groups,
            ).update(
                can_add_members_group=NamedUserGroup.objects.filter(
                    name="role:administrators",
                    realm_for_sharding=OuterRef("realm_for_sharding"),
                    is_system_group=True,
                ).values("pk")
            )

            NamedUserGroup.objects.filter(
                id__range=(lower_bound, upper_bound),
                can_add_members_group=None,
                realm_id__in=realm_ids_allowing_moderators_to_manage_all_groups,
            ).update(
                can_add_members_group=NamedUserGroup.objects.filter(
                    name="role:moderators",
                    realm_for_sharding=OuterRef("realm_for_sharding"),
                    is_system_group=True,
                ).values("pk")
            )

            # Initialize can_add_members_group to the group itself.
            # This should give the identical set of users permission,
            # since users could only exercise the permission previously
            # if they were a member of the group.
            NamedUserGroup.objects.filter(
                id__range=(lower_bound, upper_bound),
                can_add_members_group=None,
                realm_id__in=realm_ids_allowing_members_to_manage_all_groups,
            ).update(can_add_members_group=F("id"))

            # For the remaining group options, there is no direct
            # translation to can_add_members_group. We'll pick the
            # closest safe choice i.e. moderators for this.
            NamedUserGroup.objects.filter(
                id__range=(lower_bound, upper_bound),
                can_add_members_group=None,
            ).update(
                can_add_members_group=NamedUserGroup.objects.filter(
                    name="role:moderators",
                    realm_for_sharding=OuterRef("realm_for_sharding"),
                    is_system_group=True,
                ).values("pk")
            )

        lower_bound += BATCH_SIZE


class Migration(migrations.Migration):
    atomic = False
    dependencies = [
        ("zerver", "0599_namedusergroup_add_can_add_members_group"),
    ]

    operations = [
        migrations.RunPython(
            set_default_value_for_can_add_members_group,
            elidable=True,
            reverse_code=migrations.RunPython.noop,
        )
    ]
