# Copyright 2023 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 artifacts."""

import contextlib
from datetime import datetime
from io import TextIOWrapper
from typing import Optional

import django
from django.db import connection, transaction
from django.db.models import ProtectedError, Q
from django.utils import timezone

from debusine.db.models import (
    Artifact,
    ArtifactRelation,
    File,
    FileInArtifact,
    FileInStore,
)
from debusine.server.management.debusine_base_command import DebusineBaseCommand


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

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

    def add_arguments(self, parser):
        """Add CLI arguments for the delete_expired_artifacts command."""
        parser.add_argument(
            "--dry-run",
            help="Trial run: do not delete anything",
            action="store_true",
        )

    def handle(self, *args, **options):
        """Delete expired artifacts."""
        delete_expired_artifacts = DeleteExpiredArtifacts(
            self.stdout, self.stderr
        )

        delete_expired_artifacts.run(
            dry_run=options["dry_run"], verbosity=options["verbosity"]
        )

        raise SystemExit(0)


class LockTableFailed(BaseException):
    """Raised if a "LOCK TABLE" failed."""


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, out: TextIOWrapper, err: TextIOWrapper):
        """
        Initialize object.

        :param out: to write information of deleted artifacts.
        :param err: to write errors that might happen during the deletion.
        """
        self._out: TextIOWrapper = out
        self._err: TextIOWrapper = err

        self._verbosity: int = 0
        self._dry_run: bool = True

        # datetime that run() was called:
        self._initial_time: Optional[datetime] = None

        self._lock_timeout_secs: float = 5

    @staticmethod
    def _tables_to_lock():
        return [
            Artifact._meta.db_table,
            ArtifactRelation._meta.db_table,
            FileInArtifact._meta.db_table,
            FileInStore._meta.db_table,
            File._meta.db_table,
        ]

    @contextlib.contextmanager
    def lock_tables(self, table_names: list[str]):
        """
        Lock the tables or raise LockTableFailed.

        :param table_names: tables to try to lock.
        """
        with connection.cursor() as cursor, transaction.atomic():
            cursor.execute(
                f"SET LOCAL lock_timeout = '{self._lock_timeout_secs}s'"
            )

            try:
                for table_name in table_names:
                    cursor.execute(
                        f"LOCK TABLE {table_name} IN ACCESS EXCLUSIVE MODE"
                    )
            except django.db.utils.OperationalError:
                raise LockTableFailed(
                    f"Lock timed out ({self._lock_timeout_secs} seconds)"
                )

            yield cursor

    def run(self, *, dry_run: bool, verbosity: int):
        """
        Delete the expired artifacts (with dry_run or verbosity levels).

        :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._initial_time = timezone.now()  # time used for all the checks
        self._verbosity = verbosity
        self._dry_run = dry_run

        if self._dry_run and self._verbosity > 1:
            self._out.write("dry-run mode: no changes will be made\n")

        try:
            with self.lock_tables(self._tables_to_lock()):
                artifact_ids = self._mark_to_keep()

                delete_files_from_stores = self._sweep(artifact_ids)

                if self._dry_run:
                    transaction.set_rollback(True)

                    if self._verbosity > 1:
                        self._out.write("dry-run mode: no changes were made\n")

            with self.lock_tables(self._tables_to_lock()):
                # delete_files_from_stores is empty for --dry-run because
                # self._delete_artifact() does not return any files to delete
                # if self._dry_run is True

                self._delete_files_from_stores(delete_files_from_stores)
        except LockTableFailed as exc:
            self._err.write(f"{exc}. Try again.\n")

    def _mark_to_keep(self) -> set[int]:
        """
        Return artifact ids that must be kept.

        These artifacts cannot be deleted: are not expired or are expired
        but there is an ArtifactRelation such as
        target=non_expired_artifact, artifact=expired_artifact.
        """
        marked_to_keep: set[int] = set()
        visited: set[int] = set()

        assert self._initial_time is not None
        for artifact in Artifact.objects.not_expired(self._initial_time):
            self._traverse_and_mark_from_artifact(
                artifact.id, marked_to_keep, visited
            )

        return marked_to_keep

    def _delete_artifact(self, artifact: Artifact) -> set[File]:
        """
        Delete the specified Artifact and related models.

        Delete ArtifactRelation with target or artifact referencing artifact and
        the associated FileInArtifact files for this artifact.

        If the files were associated only to this artifact, it also deletes
        the database entries for the files and adds the removal of the file
        in the store for when the transaction is committed.

        :return: files that might be able to be deleted (if they don't exist
          in another artifact besides the ones that is being deleted).
        """
        if self._dry_run:
            return set()

        delete_filesobjs = set()

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

        for file_in_artifact in FileInArtifact.objects.filter(
            artifact=artifact
        ):
            fileobj = file_in_artifact.file
            file_in_artifact.delete()

            # Maybe this artifact was the only one to have this file.
            # The file needs to be deleted from the store later on,
            # after the deletion of Artifacts have been committed.

            delete_filesobjs.add(fileobj)

        artifact.delete()

        return delete_filesobjs

    def _delete_files_from_stores(self, file_objs: set[File]):
        """
        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.
        """
        if len(file_objs) > 0 and self._verbosity > 1:
            self._out.write("Deleting files from the store\n")

        for fileobj in file_objs:
            if not FileInArtifact.objects.filter(file=fileobj).exists():
                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, artifact_ids: set[int]) -> set[File]:
        """
        Delete artifacts that are not in artifact_ids and its relations.

        :return: File objects that can potentially be deleted (if not used
          by any other artifact).
        """
        assert self._initial_time is not None
        delete_files_from_stores = set()
        number_of_artifacts_deleted = 0
        with transaction.atomic():
            for artifact_expired in Artifact.objects.expired(
                self._initial_time
            ).order_by("id"):
                if (artifact_id := artifact_expired.id) not in artifact_ids:
                    number_of_artifacts_deleted += 1
                    files_to_delete = self._delete_artifact(artifact_expired)
                    delete_files_from_stores.update(files_to_delete)

                    if self._verbosity > 1:
                        self._out.write(f"Deleted artifact {artifact_id}\n")

        if self._verbosity > 1 and number_of_artifacts_deleted == 0:
            self._out.write("There were no expired artifacts\n")

        return delete_files_from_stores

    def _traverse_and_mark_from_artifact(
        self, artifact_id: int, marked_to_keep: set[int], visited: set[int]
    ):
        if artifact_id in visited:
            return

        visited.add(artifact_id)
        marked_to_keep.add(artifact_id)

        for artifact_targeting in Artifact.objects.filter(
            targeted_by__artifact_id=artifact_id
        ):
            self._traverse_and_mark_from_artifact(
                artifact_targeting.id, marked_to_keep, visited
            )
