# Copyright 2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

from functools import partial
from typing import Any

from django.db import migrations, transaction
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps


def make_data_field_renamer(
    model_name: str,
    data_field: str,
    filters: dict[str, Any],
    renames: list[tuple[str, str]],
) -> migrations.RunPython:
    """
    Make a migration operation to rename keys in data of model_name.

    This will be slow if there are many matching models.  Let's hope there
    aren't yet.
    """

    def operation(
        model_name: str,
        data_field: str,
        filters: dict[str, Any],
        renames: list[tuple[str, str]],
        apps: StateApps,
        schema_editor: BaseDatabaseSchemaEditor,
    ) -> None:
        Model = apps.get_model("db", model_name)
        filters = filters.copy()
        filters[f"{data_field}__has_any_keys"] = [key for key, _ in renames]
        while True:
            with transaction.atomic():
                models = Model.objects.select_for_update().filter(**filters)[
                    :1000
                ]
                if not models:
                    break
                for model in models:
                    data = getattr(model, data_field).copy()
                    for old_name, new_name in renames:
                        if old_name in data:
                            assert new_name not in data
                            data[new_name] = data.pop(old_name)
                    setattr(model, data_field, data)
                    model.save()

    return migrations.RunPython(
        partial(operation, model_name, data_field, filters, renames),
        reverse_code=partial(
            operation,
            model_name,
            data_field,
            filters,
            [(new_name, old_name) for old_name, new_name in renames],
        ),
    )


def make_artifact_field_renamer(
    category: str, renames: list[tuple[str, str]]
) -> migrations.RunPython:
    """Make a migration operation to rename an Artifact data field."""
    return make_data_field_renamer(
        "Artifact", "data", {"category": category}, renames
    )


def make_work_request_task_data_field_renamer(
    task_name: str, renames: list[tuple[str, str]]
) -> migrations.RunPython:
    """Make a migration operation to rename task_data in a WorkRequest."""
    return make_data_field_renamer(
        "WorkRequest", "task_data", {"task_name": task_name}, renames
    )
