# Generated by Django 5.0.6 on 2024-05-16 18:28

from collections import defaultdict

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 OuterRef, Value


def set_default_value_for_can_resolve_topics_group(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    Realm = apps.get_model("zerver", "Realm")
    NamedUserGroup = apps.get_model("zerver", "NamedUserGroup")
    UserGroup = apps.get_model("zerver", "UserGroup")
    UserGroupMembership = apps.get_model("zerver", "UserGroupMembership")
    GroupGroupMembership = apps.get_model("zerver", "GroupGroupMembership")

    Realm.objects.filter(can_resolve_topics_group=None).exclude(
        can_move_messages_between_topics_group__named_user_group=None
    ).update(
        can_resolve_topics_group=NamedUserGroup.objects.filter(
            usergroup_ptr=OuterRef("can_move_messages_between_topics_group")
        ).values("pk")
    )

    realms_with_setting_set_to_anonymous_group = Realm.objects.filter(
        can_move_messages_between_topics_group__named_user_group=None, can_resolve_topics_group=None
    )

    if len(realms_with_setting_set_to_anonymous_group) == 0:
        return

    user_groups_to_create = []
    realms_to_update = []
    anonymous_group_ids = []

    with transaction.atomic():
        for realm in realms_with_setting_set_to_anonymous_group:
            anonymous_group_ids.append(realm.can_move_messages_between_topics_group_id)

            setting_group = UserGroup(realm=realm)
            realm.can_resolve_topics_group = setting_group

            user_groups_to_create.append(setting_group)
            realms_to_update.append(realm)

        UserGroup.objects.bulk_create(user_groups_to_create)
        Realm.objects.bulk_update(realms_to_update, fields=["can_resolve_topics_group"])

        user_members = (
            UserGroupMembership.objects.filter(user_group_id__in=anonymous_group_ids)
            .annotate(member_type=Value("user"))
            .values_list("member_type", "user_group_id", "user_profile_id")
        )

        group_subgroups = (
            GroupGroupMembership.objects.filter(supergroup_id__in=anonymous_group_ids)
            .annotate(member_type=Value("group"))
            .values_list("member_type", "supergroup_id", "subgroup_id")
        )

        all_members = user_members.union(group_subgroups)

        members_dict = defaultdict(list)
        subgroups_dict = defaultdict(list)

        for member_type, group_id, member_id in all_members:
            if member_type == "user":
                members_dict[group_id].append(member_id)
            else:
                subgroups_dict[group_id].append(member_id)

        user_membership_objects = []
        group_membership_objects = []
        for realm in realms_with_setting_set_to_anonymous_group:
            members_to_add = members_dict[realm.can_move_messages_between_topics_group_id]
            subgroups_to_add = subgroups_dict[realm.can_move_messages_between_topics_group_id]

            setting_group_id = realm.can_resolve_topics_group_id
            user_membership_objects.extend(
                [
                    UserGroupMembership(user_group_id=setting_group_id, user_profile_id=user_id)
                    for user_id in members_to_add
                ]
            )
            group_membership_objects.extend(
                [
                    GroupGroupMembership(supergroup_id=setting_group_id, subgroup_id=subgroup_id)
                    for subgroup_id in subgroups_to_add
                ]
            )

        UserGroupMembership.objects.bulk_create(user_membership_objects)
        GroupGroupMembership.objects.bulk_create(group_membership_objects)


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ("zerver", "0686_realm_can_resolve_topics_group"),
    ]

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