# Copyright © 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.

"""Data migration to populate ``WorkRequest.workflow_root``."""

from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.db.models import Case, F, OuterRef, Q, Subquery, When

from debusine.artifacts.models import TaskTypes


def set_workflow_root(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    """Set ``WorkRequest.workflow_root``."""
    WorkRequest = apps.get_model("db", "WorkRequest")

    # This efficiently propagates workflow root information down from each
    # root.  The first update will update all roots, the second will update
    # all immediate children of roots, and so on.  It stops when there are
    # no more changes left to be made.
    while (
        WorkRequest.objects.filter(
            # Select work requests that are workflows or parts of workflows.
            (Q(task_type=TaskTypes.WORKFLOW) | Q(parent__isnull=False)),
            # Select work requests where workflow_root is not yet set.
            Q(workflow_root__isnull=True),
            # Select work requests that are workflow roots, or where we've
            # already set the workflow root for their parent.
            (Q(parent__isnull=True) | Q(parent__workflow_root__isnull=False)),
        ).update(
            workflow_root=Subquery(
                WorkRequest.objects.filter(id=OuterRef("id"))
                .annotate(
                    new_workflow_root=Case(
                        # If there is no parent, then this must be its own
                        # workflow root (because we selected only work
                        # requests that are workflows or have a parent).
                        When(parent__isnull=True, then=F("pk")),
                        # Otherwise, the parent's workflow root must already
                        # be set, and we can propagate that to its children.
                        default=F("parent__workflow_root"),
                    )
                )
                .values("new_workflow_root")[:1]
            )
        )
        > 0
    ):
        pass


def clear_workflow_root(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    """Clear ``WorkRequest.workflow_root``."""
    WorkRequest = apps.get_model("db", "WorkRequest")
    WorkRequest.objects.update(workflow_root=None)


class Migration(migrations.Migration):
    dependencies = [
        ("db", "0022_workrequest_workflow_root"),
    ]

    operations = [
        migrations.RunPython(set_workflow_root, clear_workflow_root),
    ]
