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

"""debusine-admin command to delete expired objects."""

import datetime as dt
import io
from typing import Any, NoReturn, override

import pgtransaction
from django.core.management import CommandParser
from django.db import transaction
from django.db.models import F, ProtectedError, Q, QuerySet
from django.template.defaultfilters import pluralize
from django.utils import timezone
from django_cte import CTE, with_cte

from debusine.db.models import (
    Artifact,
    ArtifactRelation,
    CollectionItem,
    File,
    FileInArtifact,
    FileInStore,
    Group,
    Token,
    WorkRequest,
    Workspace,
    work_requests,
)
from debusine.db.models.workspaces import DeleteWorkspaces
from debusine.django.management.debusine_base_command import DebusineBaseCommand

RETRY_ATTEMPTS = 3


class Command(DebusineBaseCommand):
    """Command to expire artifacts."""

    help = "Delete expired artifacts if no other artifact depends on them"

    @override
    def add_arguments(self, parser: CommandParser) -> None:
        parser.add_argument(
            "--dry-run",
            help="Trial run: do not delete anything",
            action="store_true",
        )

    @override
    def handle(self, *args: Any, **options: Any) -> NoReturn:
        """Delete expired artifacts."""
        with DeleteOperation(
            out=self.stdout,
            err=self.stderr,
            dry_run=options["dry_run"],
            verbosity=options["verbosity"],
        ) as operation:
            delete_expired_workrequests = DeleteExpiredWorkRequests(operation)
            delete_expired_workrequests.run()

            delete_expired_artifacts = DeleteExpiredArtifacts(operation)
            delete_expired_artifacts.run()

            delete_expired_ephemeral_groups = DeleteExpiredEphemeralGroups(
                operation
            )
            delete_expired_ephemeral_groups.run()

            delete_expired_workspaces = DeleteExpiredWorkspaces(operation)
            delete_expired_workspaces.run()

            delete_expired_tokens = DeleteExpiredTokens(operation)
            delete_expired_tokens.run()

        raise SystemExit(0)


class DeleteOperation:
    """Coordinate multiple delete operations."""

    def __init__(
        self,
        *,
        out: io.TextIOBase,
        err: io.TextIOBase,
        dry_run: bool = True,
        verbosity: int = 0,
    ) -> None:
        """
        Initialize object.

        :param out: to write information of deleted artifacts.
        :param err: to write errors that might happen during the deletion.
        :param dry_run: If True, no changes will be made to the database or
            files; instead the method will only simulate the execution
            of the actions that would have been taken.
        :param verbosity: If True, the method will provide detailed information
            about the execution progress and results.
        """
        self._out: io.TextIOBase = out
        self._err: io.TextIOBase = err

        self._verbosity: int = verbosity
        self.dry_run: bool = dry_run
        # datetime where all delete operations started
        self.initial_time: dt.datetime = (
            timezone.now()
        )  # time used for all the checks

    def __enter__(self) -> "DeleteOperation":
        """Context manager start."""
        if self.dry_run:
            self.verbose("dry-run mode: no changes will be made\n")
        return self

    def __exit__(self, *args: Any) -> None:
        """Context manager end."""
        if self.dry_run:
            self.verbose("dry-run mode: no changes were made\n")
        return None

    def verbose(self, message: str) -> None:
        """Print a verbose message."""
        if self._verbosity > 1:
            self._out.write(message)
            self._out.flush()


class DeleteExpiredArtifacts:
    """
    Delete expired artifacts.

    Delete expired artifacts that don't have a dependent artifact (direct or
    indirect via other artifacts).

    Delete other artifact/files associated models.

    Implemented using Mark and Sweep.
    """

    def __init__(self, operation: DeleteOperation) -> None:
        """
        Initialize object.

        :param operation: class used for output
        """
        self.operation = operation

    def run(self) -> None:
        """Delete the expired artifacts (with dry_run or verbosity levels)."""
        self._expire_collection_artifacts()
        self._expire_collection_metadata()
        self._expire_artifacts()
        self._delete_files_from_stores()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def _expire_collection_artifacts(self) -> None:
        """Drop artifacts from collections, retaining metadata."""
        CollectionItem.objects.drop_full_history(self.operation.initial_time)
        if self.operation.dry_run:
            transaction.set_rollback(True)

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def _expire_collection_metadata(self) -> None:
        """Delete old collection items."""
        CollectionItem.objects.drop_metadata(self.operation.initial_time)
        if self.operation.dry_run:
            transaction.set_rollback(True)

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def _expire_artifacts(self) -> None:
        """Drop old artifacts."""
        self._sweep(self._mark_to_keep())
        if self.operation.dry_run:
            transaction.set_rollback(True)

    def _mark_to_keep(
        self, *, keep_extra: QuerySet[Artifact] | None = None
    ) -> QuerySet[Artifact]:
        """
        Return a queryset of artifacts that must be kept.

        These artifacts cannot be deleted: they are not expired, or they
        would be expired but there is a relation that means they must be
        kept.

        The following artifact relations cause ``expired_artifact`` to be
        kept:

        * artifact=non_expired_artifact, target=expired_artifact, type=(any)
        * artifact=expired_artifact, target=non_expired_artifact, type=EXTENDS

        :param keep_extra: Keep these additional artifacts and anything
          marked as a result of them.
        """

        def keep_cte(cte: CTE) -> QuerySet[Artifact, Any]:
            # These artifacts are to be kept in their own right, either
            # because their expiration delays have not passed or because
            # they are part of a relevant collection.
            initial = (
                Artifact.objects.not_expired(self.operation.initial_time)
                | Artifact.objects.part_of_collection_with_retains_artifacts()
            )
            if keep_extra is not None:
                initial |= keep_extra

            # These relations cause the artifacts identified by kept_id to
            # be kept even if they would normally be expired, if keeper_id
            # is in the set of artifacts to keep.
            normal_relations = ArtifactRelation.objects.values(
                keeper_id=F("artifact"), kept_id=F("target")
            )
            extends_relations = ArtifactRelation.objects.filter(
                type=ArtifactRelation.Relations.EXTENDS
            ).values(keeper_id=F("target"), kept_id=F("artifact"))
            relations = CTE(
                normal_relations.union(extends_relations), name="relations"
            )

            # Follow relations recursively to establish the full set of
            # artifacts to keep.
            recursive = initial.values(keep_id=F("id")).union(
                with_cte(
                    relations,
                    select=cte.join(
                        relations.queryset(), keeper_id=cte.col.keep_id
                    ),
                ).values(keep_id=F("kept_id"))
            )
            assert isinstance(recursive, QuerySet)
            return recursive

        cte = CTE.recursive(keep_cte, name="keep")
        keep = with_cte(cte, select=cte.join(Artifact, id=cte.col.keep_id))
        assert isinstance(keep, QuerySet)
        return keep

    def _delete_artifact(self, artifact: Artifact) -> None:
        """
        Delete the specified Artifact and related models.

        Delete ArtifactRelation with target or artifact referencing artifact and
        the associated FileInArtifact files for this artifact.
        """
        if self.operation.dry_run:
            return

        for collection_item in artifact.collection_items.all():
            collection_item.parent_collection.manager.remove_item(
                collection_item
            )
        artifact.collection_items.update(artifact=None)

        ArtifactRelation.objects.filter(
            Q(target=artifact) | Q(artifact=artifact)
        ).delete()

        for file_in_artifact in FileInArtifact.objects.filter(
            artifact=artifact
        ):
            if hasattr(file_in_artifact, "fileupload"):
                file_in_artifact.fileupload.delete()
            file_in_artifact.delete()

        artifact.delete()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def _delete_files_from_stores(self) -> None:
        """
        Delete the files from stores.

        Files are not deleted at the same time as the artifacts. They need
        to be deleted after the transaction and only if no other FileInArtifact
        has the file.
        """
        orphaned_files = list(
            File.objects.filter(
                id__in=(
                    FileInStore.objects.values("file").difference(
                        FileInArtifact.objects.values("file")
                    )
                )
                # A SELECT ... FOR UPDATE lock prevents new references to
                # this file from being created while the lock is held.
                # There are no relations between files, so we don't need to
                # recompute anything if we need to skip some of these files
                # due to locks.
            ).select_for_update(skip_locked=True)
        )
        total = len(orphaned_files)
        if total:
            self.operation.verbose(
                f"Deleting {total} file{pluralize(total)} from the store\n"
            )

        for fileobj in orphaned_files:
            with transaction.atomic():
                # This is an orphaned file. Delete the file and
                # its dependencies
                for file_in_store in FileInStore.objects.filter(file=fileobj):
                    store = file_in_store.store.get_backend_object()
                    store.remove_file(fileobj)

                    try:
                        fileobj.delete()
                    except ProtectedError:
                        # Perhaps it's not possible to delete the file yet:
                        # it might be in another store
                        pass

    def _sweep(self, keep: QuerySet[Artifact]) -> None:
        """
        Delete artifacts that are not in ``keep``.

        This must be called within a transaction using the REPEATABLE READ
        isolation level.
        """
        number_of_artifacts_deleted = 0
        # Conceptually we should just use .exclude(id__in=keep) to find
        # expired artifacts instead, but that seems to cause PostgreSQL to
        # use a terrible plan that takes a very long time.  Pulling all the
        # IDs into Python is less elegant and will eventually be a scaling
        # problem, but with about a million IDs it only takes about four
        # seconds.
        keep_ids = set(keep.values_list("id", flat=True))
        expired = (
            Artifact.objects.expired(self.operation.initial_time)
            .exclude(id__in=keep_ids)
            .order_by("id")
        )
        expired_ids = {artifact.id for artifact in expired}
        expired_and_locked = expired.select_for_update(skip_locked=True)
        # We want this for safety checks later, but this also has the effect
        # of evaluating the queryset and therefore taking the lock.
        expired_and_locked_ids = {
            artifact.id for artifact in expired_and_locked
        }

        # Some expired artifacts may have been skipped due to a conflicting
        # lock, and some of those might be part of a relation chain that
        # should cause more artifacts to be kept.  Now that we have an upper
        # bound on the set of artifacts to expire, re-run the mark phase,
        # additionally keeping artifacts that we can't expire due to locks.
        # We're now within a FOR UPDATE lock on all the rows to be deleted,
        # so no new relations involving any of those rows can be added
        # concurrently (because that would require a FOR KEY SHARE lock on
        # the referenced rows, which conflicts with FOR UPDATE).
        unlocked_ids = expired_ids - expired_and_locked_ids
        if unlocked_ids:
            new_keep_ids = self._mark_to_keep(
                keep_extra=Artifact.objects.filter(id__in=unlocked_ids)
            )
            new_expired = (
                Artifact.objects.expired(self.operation.initial_time)
                .exclude(id__in=new_keep_ids)
                .order_by("id")
            )
            new_expired_ids = {artifact.id for artifact in new_expired}
            # Since we use the REPEATABLE READ isolation level,
            # nonrepeatable reads and (in PostgreSQL) phantom reads are
            # impossible, so _mark_to_keep must keep at least all the
            # rows it returned the first time round.
            if not (new_expired_ids <= expired_and_locked_ids):
                raise AssertionError(
                    "Second expired artifacts mark phase wants to expire"
                    " some artifacts that were not expired in the first"
                    " phase, despite REPEATABLE READ!"
                )
            expired_and_locked = new_expired

        for artifact_expired in expired_and_locked:
            artifact_id = artifact_expired.id
            number_of_artifacts_deleted += 1
            self._delete_artifact(artifact_expired)

            self.operation.verbose(f"Deleted artifact {artifact_id}\n")

        if number_of_artifacts_deleted == 0:
            self.operation.verbose("There were no expired artifacts\n")


class DeleteExpiredWorkRequests:
    """
    Delete expired work requests.

    Delete work requests that don't have unexpired child work requests.

    Delete also the work requests' internal collection if present.
    """

    to_delete: work_requests.DeleteUnused

    def __init__(self, operation: DeleteOperation) -> None:
        """
        Initialize object.

        :param operation: class used for output
        """
        self.operation = operation

    def perform_deletions(self) -> None:
        """Delete objects found in self.work_requests and self.collections."""
        self.to_delete.perform_deletions()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def run(self) -> None:
        """Find and optionally delete the expired work requests."""
        self.to_delete = work_requests.DeleteUnused(
            WorkRequest.objects.expired(at=self.operation.initial_time)
        )
        if self.to_delete.work_requests:
            self.operation.verbose(
                f"Deleting {len(self.to_delete.work_requests)}"
                " expired work requests"
                f" and {len(self.to_delete.collections)}"
                " expired collections\n"
            )

            # FIXME: do we enumerate each deleted work request to stdout?

            if not self.operation.dry_run:
                self.perform_deletions()
        else:
            self.operation.verbose("There were no expired work requests\n")


class DeleteExpiredEphemeralGroups:
    """Delete expired exphemeral groups."""

    groups: QuerySet[Group]

    def __init__(self, operation: DeleteOperation) -> None:
        """
        Initialize object.

        :param operation: class used for output
        """
        self.operation = operation

    def perform_deletions(self) -> None:
        """Delete previously found unused ephemeral groups."""
        self.groups.delete()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def run(self) -> None:
        """Find and optionally delete the expired work requests."""
        self.groups = Group.objects.unused_ephemeral()
        if count := len(self.groups):
            self.operation.verbose(
                f"Deleting {count} unused ephemeral groups\n"
            )
            if not self.operation.dry_run:
                self.perform_deletions()
        else:
            self.operation.verbose("There were no unused ephemeral groups\n")


class DeleteExpiredWorkspaces:
    """
    Delete expired workspaces.

    Delete workspaces that have expiration_delay set, ``expiration_delay`` days
    after ``max(created_at, last_task_completion_time)``
    """

    #: Work requests to delete
    to_delete: DeleteWorkspaces

    def __init__(self, operation: DeleteOperation) -> None:
        """
        Initialize object.

        :param operation: class used for output
        """
        self.operation = operation

    def perform_deletions(self) -> None:
        """Delete objects found in self.work_requests and self.collections."""
        self.to_delete.perform_deletions()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def run(self) -> None:
        """Find and optionally delete the expired work requests."""
        self.to_delete = DeleteWorkspaces(
            Workspace.objects.exclude(expiration_delay__isnull=True)
            .with_expiration_time()
            .filter(expiration_time__lt=self.operation.initial_time)
        )
        if self.to_delete.workspaces:
            self.operation.verbose(
                f"Deleting {len(self.to_delete.workspaces)}"
                " expired workspaces\n"
            )
            if not self.operation.dry_run:
                self.perform_deletions()
        else:
            self.operation.verbose("There were no expired workspaces\n")


class DeleteExpiredTokens:
    """Delete expired tokens."""

    tokens: QuerySet[Token]

    def __init__(self, operation: DeleteOperation) -> None:
        """
        Initialize object.

        :param operation: class used for output
        """
        self.operation = operation

    def perform_deletions(self) -> None:
        """Delete previously-found expired tokens."""
        self.tokens.delete()

    @pgtransaction.atomic(
        isolation_level=pgtransaction.REPEATABLE_READ, retry=RETRY_ATTEMPTS
    )
    def run(self) -> None:
        """Find and optionally delete expired tokens."""
        self.tokens = (
            Token.objects.expired(at=self.operation.initial_time)
            # Expiring tokens linked to workers would violate database
            # constraints.  The provisioner should clean up the most likely
            # case of this (a worker that never finished activating), after
            # which we can expire such tokens.
            .exclude(worker__isnull=False)
            .exclude(activating_worker__isnull=False)
        )
        if count := len(self.tokens):
            self.operation.verbose(f"Deleting {count} expired tokens\n")
            if not self.operation.dry_run:
                self.perform_deletions()
        else:
            self.operation.verbose("There were no expired tokens\n")
