import collections
import logging
import os
import shutil
from datetime import datetime, timedelta, timezone
from difflib import unified_diff
from typing import Any, TypeAlias

import orjson
import pyvips
from bs4 import BeautifulSoup
from django.conf import settings
from django.core.management.base import CommandError
from django.core.validators import validate_email
from django.db import connection, transaction
from django.db.backends.utils import CursorWrapper
from django.db.models import Count, Q
from django.utils.timezone import now as timezone_now
from psycopg2.extras import execute_values
from psycopg2.sql import SQL, Identifier

from analytics.models import RealmCount, StreamCount, UserCount
from version import ZULIP_VERSION
from zerver.actions.create_realm import set_default_for_realm_permission_group_settings
from zerver.actions.realm_settings import (
    do_change_realm_plan_type,
    do_set_realm_new_stream_announcements_stream,
    do_set_realm_zulip_update_announcements_stream,
)
from zerver.actions.user_settings import do_change_avatar_fields
from zerver.lib.avatar_hash import user_avatar_base_path_from_ids
from zerver.lib.bulk_create import bulk_set_users_or_streams_recipient_fields
from zerver.lib.export import DATE_FIELDS, Field, Path, Record, TableName
from zerver.lib.markdown import markdown_convert
from zerver.lib.markdown import version as markdown_version
from zerver.lib.message import get_last_message_id
from zerver.lib.migration_status import MigrationStatusJson, parse_migration_status
from zerver.lib.mime_types import guess_type
from zerver.lib.onboarding import (
    OnboardingMessageTypeEnum,
    send_initial_direct_messages_to_user,
    send_initial_realm_messages,
)
from zerver.lib.parallel import run_parallel
from zerver.lib.partial import partial
from zerver.lib.push_notifications import sends_notifications_directly
from zerver.lib.remote_server import maybe_enqueue_audit_log_upload
from zerver.lib.server_initialization import create_internal_realm, server_initialized
from zerver.lib.streams import (
    get_stream_permission_default_group,
    render_stream_description,
    update_stream_active_status_for_realm,
)
from zerver.lib.thumbnail import (
    THUMBNAIL_ACCEPT_IMAGE_TYPES,
    BadImageError,
    manifest_and_get_user_upload_previews,
    maybe_thumbnail,
)
from zerver.lib.timestamp import datetime_to_timestamp
from zerver.lib.upload import ensure_avatar_image, sanitize_name, upload_backend, upload_emoji_image
from zerver.lib.upload.s3 import get_bucket
from zerver.lib.user_counts import realm_user_count_by_role
from zerver.lib.user_groups import create_system_user_groups_for_realm
from zerver.lib.user_message import UserMessageLite, bulk_insert_ums
from zerver.lib.utils import generate_api_key, process_list_in_batches
from zerver.lib.zulip_update_announcements import send_zulip_update_announcements_to_realm
from zerver.models import (
    AlertWord,
    Attachment,
    BotConfigData,
    BotStorageData,
    ChannelFolder,
    Client,
    CustomProfileField,
    CustomProfileFieldValue,
    DefaultStream,
    DirectMessageGroup,
    GroupGroupMembership,
    Message,
    MutedUser,
    NamedUserGroup,
    NavigationView,
    OnboardingStep,
    OnboardingUserMessage,
    Reaction,
    Realm,
    RealmAuditLog,
    RealmAuthenticationMethod,
    RealmDomain,
    RealmEmoji,
    RealmFilter,
    RealmPlayground,
    RealmUserDefault,
    Recipient,
    SavedSnippet,
    ScheduledMessage,
    Service,
    Stream,
    Subscription,
    UserActivity,
    UserActivityInterval,
    UserGroup,
    UserGroupMembership,
    UserMessage,
    UserPresence,
    UserProfile,
    UserStatus,
    UserTopic,
)
from zerver.models.groups import SystemGroups
from zerver.models.messages import SubMessage
from zerver.models.presence import PresenceSequence
from zerver.models.realm_audit_logs import AuditLogEventType
from zerver.models.realms import get_realm
from zerver.models.recipients import get_direct_message_group_hash
from zerver.models.users import get_system_bot, get_user_profile_by_id
from zproject.backends import AUTH_BACKEND_NAME_MAP

ImportedTableData: TypeAlias = dict[str, list[Record]]

realm_tables = [
    ("zerver_realmauthenticationmethod", RealmAuthenticationMethod, "realmauthenticationmethod"),
    ("zerver_defaultstream", DefaultStream, "defaultstream"),
    ("zerver_realmemoji", RealmEmoji, "realmemoji"),
    ("zerver_realmdomain", RealmDomain, "realmdomain"),
    ("zerver_realmfilter", RealmFilter, "realmfilter"),
    ("zerver_realmplayground", RealmPlayground, "realmplayground"),
]  # List[Tuple[TableName, Any, str]]


# ID_MAP is a dictionary that maps table names to dictionaries
# that map old ids to new ids.  We use this in
# re_map_foreign_keys and other places.
#
# We explicitly initialize ID_MAP with the tables that support
# id re-mapping.
#
# Code reviewers: give these tables extra scrutiny, as we need to
# make sure to reload related tables AFTER we re-map the ids.
ID_MAP: dict[str, dict[int, int]] = {
    "alertword": {},
    "client": {},
    "user_profile": {},
    "huddle": {},
    "realm": {},
    "stream": {},
    "recipient": {},
    "subscription": {},
    "defaultstream": {},
    "onboardingstep": {},
    "presencesequence": {},
    "reaction": {},
    "realmauthenticationmethod": {},
    "realmemoji": {},
    "realmdomain": {},
    "realmfilter": {},
    "realmplayground": {},
    "message": {},
    "user_presence": {},
    "userstatus": {},
    "useractivity": {},
    "useractivityinterval": {},
    "usermessage": {},
    "customprofilefield": {},
    "customprofilefieldvalue": {},
    "attachment": {},
    "realmauditlog": {},
    "recipient_to_huddle_map": {},
    "usertopic": {},
    "muteduser": {},
    "service": {},
    "usergroup": {},
    "usergroupmembership": {},
    "groupgroupmembership": {},
    "botstoragedata": {},
    "botconfigdata": {},
    "analytics_realmcount": {},
    "analytics_streamcount": {},
    "analytics_usercount": {},
    "realmuserdefault": {},
    "scheduledmessage": {},
    "onboardingusermessage": {},
    "savedsnippet": {},
    "channelfolder": {},
    "navigationview": {},
    "submessage": {},
}

id_map_to_list: dict[str, dict[int, list[int]]] = {
    "huddle_to_user_list": {},
}

path_maps: dict[str, dict[str, str]] = {
    # Maps original attachment path pre-import to the final, post-import
    # attachment path.
    "old_attachment_path_to_new_path": {},
    # Inverse of old_attachment_path_to_new_path.
    "new_attachment_path_to_old_path": {},
    # Maps the new (post-import) attachment path to the absolute path to the file
    # in the on-disk export data that we're importing.
    # Allows code running after this is filled to access file contents
    # without needing to go through S3 to get it.
    "new_attachment_path_to_local_data_path": {},
}

message_id_to_attachments: dict[str, dict[int, list[str]]] = {
    "zerver_message": collections.defaultdict(list),
    "zerver_scheduledmessage": collections.defaultdict(list),
}


def map_messages_to_attachments(data: ImportedTableData) -> None:
    for attachment in data["zerver_attachment"]:
        for message_id in attachment["messages"]:
            message_id_to_attachments["zerver_message"][message_id].append(attachment["path_id"])

        for scheduled_message_id in attachment["scheduled_messages"]:
            message_id_to_attachments["zerver_scheduledmessage"][scheduled_message_id].append(
                attachment["path_id"]
            )


def update_id_map(table: TableName, old_id: int, new_id: int) -> None:
    if table not in ID_MAP:
        raise Exception(
            f"""
            Table {table} is not initialized in ID_MAP, which could
            mean that we have not thought through circular
            dependencies.
            """
        )
    ID_MAP[table][old_id] = new_id


def fix_datetime_fields(data: ImportedTableData, table: TableName) -> None:
    for item in data[table]:
        for field_name in DATE_FIELDS[table]:
            if item[field_name] is not None:
                item[field_name] = datetime.fromtimestamp(item[field_name], tz=timezone.utc)


def fix_upload_links(data: ImportedTableData, message_table: TableName) -> None:
    """
    Because the URLs for uploaded files encode the realm ID of the
    organization being imported (which is only determined at import
    time), we need to rewrite the URLs of links to uploaded files
    during the import process.

    Applied to attachments path_id found in messages of zerver_message and zerver_scheduledmessage tables.
    """
    for message in data[message_table]:
        if message["has_attachment"] is True:
            for attachment_path in message_id_to_attachments[message_table][message["id"]]:
                old_path = path_maps["new_attachment_path_to_old_path"][attachment_path]
                message["content"] = message["content"].replace(old_path, attachment_path)

                if message["rendered_content"]:
                    message["rendered_content"] = message["rendered_content"].replace(
                        old_path, attachment_path
                    )


def fix_stream_permission_group_settings(
    data: ImportedTableData, system_groups_name_dict: dict[str, NamedUserGroup]
) -> None:
    table = get_db_table(Stream)
    for stream in data[table]:
        for setting_name in Stream.stream_permission_group_settings:
            if setting_name == "can_send_message_group" and "stream_post_policy" in stream:
                # This block is needed when importing data for Rocket.Chat.
                # While converting the data into Zulip supported format,
                # stream_post_policy for read-only Rocket.Chat channel is
                # set to STREAM_POST_POLICY_MODERATORS and so we need to
                # set the can_send_message_group setting accordingly.
                if stream["stream_post_policy"] == Stream.STREAM_POST_POLICY_MODERATORS:
                    stream[setting_name] = system_groups_name_dict[SystemGroups.MODERATORS]
                else:
                    stream[setting_name] = get_stream_permission_default_group(
                        setting_name, system_groups_name_dict
                    )

                # Delete the stream_post_policy once can_send_message_group
                # field is set correctly.
                del stream["stream_post_policy"]

                continue

            stream[setting_name] = get_stream_permission_default_group(
                setting_name, system_groups_name_dict
            )


def create_subscription_events(data: ImportedTableData, realm_id: int) -> None:
    """
    When the export data doesn't contain the table `zerver_realmauditlog`,
    this function creates RealmAuditLog objects for `subscription_created`
    type event for all the existing Stream subscriptions.

    This is needed for all the export tools which do not include the
    table `zerver_realmauditlog` (e.g. Slack) because the appropriate
    data about when a user was subscribed is not exported by the third-party
    service.
    """
    all_subscription_logs = []

    event_last_message_id = get_last_message_id()
    event_time = timezone_now()

    recipient_id_to_stream_id = {
        d["id"]: d["type_id"] for d in data["zerver_recipient"] if d["type"] == Recipient.STREAM
    }

    for sub in data["zerver_subscription"]:
        recipient_id = sub["recipient_id"]
        stream_id = recipient_id_to_stream_id.get(recipient_id)

        if stream_id is None:
            continue

        user_id = sub["user_profile_id"]

        all_subscription_logs.append(
            RealmAuditLog(
                realm_id=realm_id,
                acting_user_id=user_id,
                modified_user_id=user_id,
                modified_stream_id=stream_id,
                event_last_message_id=event_last_message_id,
                event_time=event_time,
                event_type=AuditLogEventType.SUBSCRIPTION_CREATED,
            )
        )
    RealmAuditLog.objects.bulk_create(all_subscription_logs)


def fix_service_tokens(data: ImportedTableData, table: TableName) -> None:
    """
    The tokens in the services are created by 'generate_api_key'.
    As the tokens are unique, they should be re-created for the imports.
    """
    for item in data[table]:
        item["token"] = generate_api_key()


def process_direct_message_group_hash(data: ImportedTableData, table: TableName) -> None:
    """
    Build new direct message group hashes with the updated ids of the users
    """
    for direct_message_group in data[table]:
        user_id_list = id_map_to_list["huddle_to_user_list"][direct_message_group["id"]]
        direct_message_group["huddle_hash"] = get_direct_message_group_hash(user_id_list)


def get_direct_message_groups_from_subscription(data: ImportedTableData, table: TableName) -> None:
    """
    Extract the IDs of the user_profiles involved in a direct message group from
    the subscription object
    This helps to generate a unique direct message group hash from the updated
    user_profile ids
    """
    id_map_to_list["huddle_to_user_list"] = {
        value: [] for value in ID_MAP["recipient_to_huddle_map"].values()
    }

    for subscription in data[table]:
        if subscription["recipient"] in ID_MAP["recipient_to_huddle_map"]:
            direct_message_group_id = ID_MAP["recipient_to_huddle_map"][subscription["recipient"]]
            id_map_to_list["huddle_to_user_list"][direct_message_group_id].append(
                subscription["user_profile_id"]
            )


def fix_customprofilefield(data: ImportedTableData) -> None:
    """
    In CustomProfileField with 'field_type' like 'USER', the IDs need to be
    re-mapped.
    """
    field_type_USER_ids = {
        item["id"]
        for item in data["zerver_customprofilefield"]
        if item["field_type"] == CustomProfileField.USER
    }

    for item in data["zerver_customprofilefieldvalue"]:
        if item["field_id"] in field_type_USER_ids:
            old_user_id_list = orjson.loads(item["value"])

            new_id_list = re_map_foreign_keys_many_to_many_internal(
                table="zerver_customprofilefieldvalue",
                field_name="value",
                related_table="user_profile",
                old_id_list=old_user_id_list,
            )
            item["value"] = orjson.dumps(new_id_list).decode()


def fix_message_rendered_content(
    realm: Realm,
    sender_map: dict[int, Record],
    messages: list[Record],
    content_key: str = "content",
    rendered_content_key: str = "rendered_content",
) -> None:
    """
    This function sets the rendered_content of the messages we're importing.
    """
    for message in messages:
        if content_key not in message:
            # Message-edit entries include topic moves, which don't
            # have any content changes to process.
            continue

        if message[rendered_content_key] is not None:
            # For Zulip->Zulip imports, we use the original rendered
            # Markdown; this avoids issues where e.g. a mention can no
            # longer render properly because a user has changed their
            # name.
            #
            # However, we still need to update the data-user-id and
            # similar values stored on mentions, stream mentions, and
            # similar syntax in the rendered HTML.
            soup = BeautifulSoup(message[rendered_content_key], "html.parser")

            user_mentions = soup.find_all("span", {"class": "user-mention"})
            if len(user_mentions) != 0:
                user_id_map = ID_MAP["user_profile"]
                for mention in user_mentions:
                    if not mention.has_attr("data-user-id"):
                        # Legacy mentions don't have a data-user-id
                        # field; we should just import them
                        # unmodified.
                        continue
                    if mention["data-user-id"] == "*":
                        # No rewriting is required for wildcard mentions
                        continue
                    old_user_id = int(mention["data-user-id"])
                    if old_user_id in user_id_map:
                        mention["data-user-id"] = str(user_id_map[old_user_id])
                message[rendered_content_key] = str(soup)

            stream_mentions = soup.find_all("a", {"class": "stream"})
            if len(stream_mentions) != 0:
                stream_id_map = ID_MAP["stream"]
                for mention in stream_mentions:
                    old_stream_id = int(mention["data-stream-id"])
                    if old_stream_id in stream_id_map:
                        mention["data-stream-id"] = str(stream_id_map[old_stream_id])
                message[rendered_content_key] = str(soup)

            user_group_mentions = soup.find_all("span", {"class": "user-group-mention"})
            if len(user_group_mentions) != 0:
                user_group_id_map = ID_MAP["usergroup"]
                for mention in user_group_mentions:
                    old_user_group_id = int(mention["data-user-group-id"])
                    if old_user_group_id in user_group_id_map:
                        mention["data-user-group-id"] = str(user_group_id_map[old_user_group_id])
                message[rendered_content_key] = str(soup)

            # Enqueue thumbnailing of all thumbnails in the message.
            # We do this without taking a lock on the ImageAttachment
            # rows, because the race here is that the images
            # referenced by this message were encountered previously,
            # and are currently being thumbnailed -- which will
            # worst-case result in the images being enqueued a second
            # time, and a no-op when those are processed.  The return
            # value will also be out of date -- but that is irrelevant
            # in this use case.
            manifest_and_get_user_upload_previews(realm.id, message[content_key])

            continue

        try:
            content = message[content_key]

            sender_id = message["sender_id"]
            sender = sender_map[sender_id]
            sent_by_bot = sender["is_bot"]
            translate_emoticons = sender["translate_emoticons"]

            # We don't handle alert words on import from third-party
            # platforms, since they generally don't have an "alert
            # words" type feature, and notifications aren't important anyway.
            realm_alert_words_automaton = None

            # This also enqueues thumbnailing for images that are referenced
            rendered_content = markdown_convert(
                content=content,
                realm_alert_words_automaton=realm_alert_words_automaton,
                message_realm=realm,
                sent_by_bot=sent_by_bot,
                translate_emoticons=translate_emoticons,
            ).rendered_content

            message[rendered_content_key] = rendered_content
            if "scheduled_timestamp" not in message:
                # This logic runs also for ScheduledMessage, which doesn't use
                # the rendered_content_version field.
                message["rendered_content_version"] = markdown_version
        except Exception:
            # This generally happens with two possible causes:
            # * rendering Markdown throwing an uncaught exception
            # * rendering Markdown failing with the exception being
            #   caught in Markdown (which then returns None, causing the
            #   rendered_content assert above to fire).
            logging.warning(
                "Error in Markdown rendering for message ID %s; continuing", message["id"]
            )


def fix_message_edit_history(
    realm: Realm, sender_map: dict[int, Record], messages: list[Record]
) -> None:
    user_id_map = ID_MAP["user_profile"]
    for message in messages:
        edit_history_json = message.get("edit_history")
        if not edit_history_json:
            continue

        edit_history = orjson.loads(edit_history_json)
        for edit_history_message_dict in edit_history:
            edit_history_message_dict["user_id"] = user_id_map[edit_history_message_dict["user_id"]]

        fix_message_rendered_content(
            realm,
            sender_map,
            messages=edit_history,
            content_key="prev_content",
            rendered_content_key="prev_rendered_content",
        )

        message["edit_history"] = orjson.dumps(edit_history).decode()


def current_table_ids(data: ImportedTableData, table: TableName) -> list[int]:
    """
    Returns the ids present in the current table
    """
    return [item["id"] for item in data[table]]


def idseq(model_class: Any, cursor: CursorWrapper) -> str:
    sequences = connection.introspection.get_sequences(cursor, model_class._meta.db_table)
    for sequence in sequences:
        if sequence["column"] == "id":
            return sequence["name"]
    raise Exception(f"No sequence found for 'id' of {model_class}")


def allocate_ids(model_class: Any, count: int) -> list[int]:
    """
    Increases the sequence number for a given table by the amount of objects being
    imported into that table. Hence, this gives a reserved range of IDs to import the
    converted Slack objects into the tables.
    """
    with connection.cursor() as cursor:
        sequence = idseq(model_class, cursor)
        cursor.execute("select nextval(%s) from generate_series(1, %s)", [sequence, count])
        query = cursor.fetchall()  # Each element in the result is a tuple like (5,)

    # convert List[Tuple[int]] to List[int]
    return [item[0] for item in query]


def convert_to_id_fields(data: ImportedTableData, table: TableName, field_name: Field) -> None:
    """
    When Django gives us dict objects via model_to_dict, the foreign
    key fields are `foo`, but we want `foo_id` for the bulk insert.
    This function handles the simple case where we simply rename
    the fields.  For cases where we need to munge ids in the
    database, see re_map_foreign_keys.
    """
    for item in data[table]:
        item[field_name + "_id"] = item[field_name]
        del item[field_name]


def re_map_foreign_keys(
    data: ImportedTableData,
    table: TableName,
    field_name: Field,
    related_table: TableName,
    verbose: bool = False,
    id_field: bool = False,
    recipient_field: bool = False,
) -> None:
    """
    This is a wrapper function for all the realm data tables
    and only avatar and attachment records need to be passed through the internal function
    because of the difference in data format (ImportedTableData corresponding to realm data tables
    and List[Record] corresponding to the avatar and attachment records)
    """

    # See comments in bulk_import_user_message_data.
    assert related_table != "usermessage"

    re_map_foreign_keys_internal(
        data[table],
        table,
        field_name,
        related_table,
        verbose,
        id_field,
        recipient_field,
    )


def re_map_foreign_keys_internal(
    data_table: list[Record],
    table: TableName,
    field_name: Field,
    related_table: TableName,
    verbose: bool = False,
    id_field: bool = False,
    recipient_field: bool = False,
) -> None:
    """
    We occasionally need to assign new ids to rows during the
    import/export process, to accommodate things like existing rows
    already being in tables.  See bulk_import_client for more context.

    The tricky part is making sure that foreign key references
    are in sync with the new ids, and this fixer function does
    the re-mapping.  (It also appends `_id` to the field.)
    """
    lookup_table = ID_MAP[related_table]
    for item in data_table:
        old_id = item[field_name]
        if recipient_field:
            if related_table == "stream" and item["type"] == 2:
                pass
            elif related_table == "user_profile" and item["type"] == 1:
                pass
            elif related_table == "huddle" and item["type"] == 3:
                # save the recipient id with the direct message group id, so that
                # we can extract the user_profile ids involved in a direct message
                # group with the help of the subscription object
                # check function 'get_direct_message_groups_from_subscription'
                ID_MAP["recipient_to_huddle_map"][item["id"]] = lookup_table[old_id]
            else:
                continue
        old_id = item[field_name]
        if old_id in lookup_table:
            new_id = lookup_table[old_id]
            if verbose:
                logging.info(
                    "Remapping %s %s from %s to %s", table, field_name + "_id", old_id, new_id
                )
        else:
            new_id = old_id
        if not id_field:
            item[field_name + "_id"] = new_id
            del item[field_name]
        else:
            item[field_name] = new_id


def re_map_realm_emoji_codes(data: ImportedTableData, *, table_name: str) -> None:
    """
    Some tables, including Reaction and UserStatus, contain a form of
    foreign key reference to the RealmEmoji table in the form of
    `str(realm_emoji.id)` when `reaction_type="realm_emoji"`.

    See the block comment for emoji_code in the AbstractEmoji
    definition for more details.
    """
    realm_emoji_dct = {}

    for row in data["zerver_realmemoji"]:
        realm_emoji_dct[row["id"]] = row

    for row in data[table_name]:
        if row["reaction_type"] == Reaction.REALM_EMOJI:
            old_realm_emoji_id = int(row["emoji_code"])

            # Fail hard here if we didn't map correctly here
            new_realm_emoji_id = ID_MAP["realmemoji"][old_realm_emoji_id]

            # This is a very important sanity check.
            realm_emoji_row = realm_emoji_dct[new_realm_emoji_id]
            assert realm_emoji_row["name"] == row["emoji_name"]

            # Now update emoji_code to the new id.
            row["emoji_code"] = str(new_realm_emoji_id)


def re_map_foreign_keys_many_to_many(
    data: ImportedTableData,
    table: TableName,
    field_name: Field,
    related_table: TableName,
    verbose: bool = False,
) -> None:
    """
    We need to assign new ids to rows during the import/export
    process.

    The tricky part is making sure that foreign key references
    are in sync with the new ids, and this wrapper function does
    the re-mapping only for ManyToMany fields.
    """
    for item in data[table]:
        old_id_list = item[field_name]
        new_id_list = re_map_foreign_keys_many_to_many_internal(
            table, field_name, related_table, old_id_list, verbose
        )
        item[field_name] = new_id_list
        del item[field_name]


def re_map_foreign_keys_many_to_many_internal(
    table: TableName,
    field_name: Field,
    related_table: TableName,
    old_id_list: list[int],
    verbose: bool = False,
) -> list[int]:
    """
    This is an internal function for tables with ManyToMany fields,
    which takes the old ID list of the ManyToMany relation and returns the
    new updated ID list.
    """
    lookup_table = ID_MAP[related_table]
    new_id_list = []
    for old_id in old_id_list:
        if old_id in lookup_table:
            new_id = lookup_table[old_id]
            if verbose:
                logging.info(
                    "Remapping %s %s from %s to %s", table, field_name + "_id", old_id, new_id
                )
        else:
            new_id = old_id
        new_id_list.append(new_id)
    return new_id_list


def fix_bitfield_keys(data: ImportedTableData, table: TableName, field_name: Field) -> None:
    for item in data[table]:
        item[field_name] = item[field_name + "_mask"]
        del item[field_name + "_mask"]


def remove_denormalized_recipient_column_from_data(data: ImportedTableData) -> None:
    """
    The recipient column shouldn't be imported, we'll set the correct values
    when Recipient table gets imported.
    """
    for stream_dict in data["zerver_stream"]:
        if "recipient" in stream_dict:
            del stream_dict["recipient"]

    for user_profile_dict in data["zerver_userprofile"]:
        if "recipient" in user_profile_dict:
            del user_profile_dict["recipient"]

    for direct_message_group_dict in data["zerver_huddle"]:
        if "recipient" in direct_message_group_dict:
            del direct_message_group_dict["recipient"]


def get_db_table(model_class: Any) -> str:
    """E.g. (RealmDomain -> 'zerver_realmdomain')"""
    return model_class._meta.db_table


def update_model_ids(model: Any, data: ImportedTableData, related_table: TableName) -> None:
    table = get_db_table(model)

    # Important: remapping usermessage rows is
    # not only unnecessary, it's expensive and can cause
    # memory errors. We don't even use ids from ID_MAP.
    assert table != "usermessage"

    old_id_list = current_table_ids(data, table)
    allocated_id_list = allocate_ids(model, len(data[table]))
    for item in range(len(data[table])):
        update_id_map(related_table, old_id_list[item], allocated_id_list[item])
    re_map_foreign_keys(data, table, "id", related_table=related_table, id_field=True)


def bulk_import_user_message_data(data: ImportedTableData, dump_file_id: int) -> None:
    model = UserMessage
    table = "zerver_usermessage"
    lst = data[table]

    # IMPORTANT NOTE: We do not use any primary id
    # data from either the import itself or ID_MAP.
    # We let the DB itself generate ids.  Note that
    # no tables use user_message.id as a foreign key,
    # so we can safely avoid all re-mapping complexity.

    def process_batch(items: list[dict[str, Any]]) -> None:
        ums = [
            UserMessageLite(
                user_profile_id=item["user_profile_id"],
                message_id=item["message_id"],
                flags=item["flags"],
            )
            for item in items
        ]
        bulk_insert_ums(ums)

    chunk_size = 10000

    process_list_in_batches(
        lst=lst,
        chunk_size=chunk_size,
        process_batch=process_batch,
    )

    logging.info("Successfully imported %s from %s[%s].", model, table, dump_file_id)


def bulk_import_model(data: ImportedTableData, model: Any, dump_file_id: str | None = None) -> None:
    table = get_db_table(model)
    # TODO, deprecate dump_file_id
    model.objects.bulk_create(model(**item) for item in data[table])
    if dump_file_id is None:
        logging.info("Successfully imported %s from %s.", model, table)
    else:
        logging.info("Successfully imported %s from %s[%s].", model, table, dump_file_id)


def bulk_import_named_user_groups(data: ImportedTableData) -> None:
    vals = [
        (
            group["usergroup_ptr_id"],
            group["realm_for_sharding_id"],
            group["name"],
            group["description"],
            group["is_system_group"],
            group["can_add_members_group_id"],
            group["can_join_group_id"],
            group["can_leave_group_id"],
            group["can_manage_group_id"],
            group["can_mention_group_id"],
            group["can_remove_members_group_id"],
            group["deactivated"],
            group["date_created"],
        )
        for group in data["zerver_namedusergroup"]
    ]

    query = SQL(
        """
        INSERT INTO zerver_namedusergroup (usergroup_ptr_id, realm_id, name, description, is_system_group, can_add_members_group_id, can_join_group_id,  can_leave_group_id, can_manage_group_id, can_mention_group_id, can_remove_members_group_id, deactivated, date_created)
        VALUES %s
        """
    )
    with connection.cursor() as cursor:
        execute_values(cursor.cursor, query, vals)


# Client is a table shared by multiple realms, so in order to
# correctly import multiple realms into the same server, we need to
# check if a Client object already exists, and so we need to support
# remap all Client IDs to the values in the new DB.
def bulk_import_client(data: ImportedTableData, model: Any, table: TableName) -> None:
    for item in data[table]:
        try:
            client = Client.objects.get(name=item["name"])
        except Client.DoesNotExist:
            client = Client.objects.create(name=item["name"])
        update_id_map(table="client", old_id=item["id"], new_id=client.id)


def set_subscriber_count_for_channels(realm: Realm) -> None:
    query = """
        UPDATE zerver_stream
        SET subscriber_count = (
            SELECT COUNT(*)
