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

"""Tar-related classes: generate tar.gz for artifacts, etc."""

import logging
import os
import shutil
import tarfile
import tempfile
from collections.abc import Iterator
from pathlib import Path

from debusine.db.models import Artifact, File


class TarArtifact:
    """Create a .tar.gz file."""

    def __init__(
        self, artifact: Artifact, subdirectory: str | None = None
    ) -> None:
        """Initialize member variables."""
        self._artifact = artifact

        self._files = self._get_file_list(artifact, subdirectory)
        self._pos = 0
        self._max_chunk_size = 50 * 1024 * 1024

        self._temp_directory = Path(
            tempfile.mkdtemp(prefix=f"debusine-artifact-download-{artifact.id}")
        )

        tar_file = self._temp_directory / f"artifact-{artifact.id}.tar.gz"

        self._tar = tarfile.open(tar_file.as_posix(), mode="w|gz")  # noqa: SIM115
        self._tar_closed = False

        self._tar_file = tar_file.open("rb")

        self._block_size = os.fstat(self._tar_file.fileno()).st_blksize
        self._bytes_deleted = 0

        self._closed = False

    @staticmethod
    def _get_file_list(
        artifact: Artifact, subdirectory: str | None
    ) -> dict[str, File]:
        """Return dictionary of file paths - file objects to be included."""
        files = {}

        if subdirectory is None:
            subdirectory = ""

        for file_in_artifact in (
            artifact.fileinartifact_set.select_related("file")
            .filter(path__startswith=subdirectory, complete=True)
            .order_by("-path")
        ):
            files[file_in_artifact.path] = file_in_artifact.file

        return files

    def __iter__(self) -> Iterator[bytes]:
        """Self is an iterator."""
        return self

    def __next__(self) -> bytes:
        """
        Return data or raise StopIteration if any exception happened.

        StopIteration if any exception occurred is needed to make debusine
        server to close the connection with the client and to avoid the
        client waiting for new data when no new data will be generated.

        The exception is logged for debugging purposes.
        """
        try:
            return self.next()
        except Exception as exc:
            if not isinstance(exc, StopIteration):
                logging.exception(exc)  # noqa: G200
            self.close()
            raise StopIteration

    def flush_tar(self) -> None:
        """Close the tar file so that the last pending data gets written."""
        self._tar.close()
        self._tar_closed = True

    def next(self) -> bytes:
        """Return pending data (from the last added file) and adds a file."""
        data = self._tar_file.read(self._max_chunk_size)

        if data != b"":
            # No need to add files, close, etc. if data can be
            # returned
            return data

        # No data read: add files, close Tar, StopIterating, etc.
        if len(self._files) > 0:
            # No data read and there are new files to add
            path, fileobj = self._files.popitem()

            file_backend = self._artifact.workspace.scope.download_file_backend(
                fileobj
            )
            with file_backend.get_stream(fileobj) as file:
                tarinfo = tarfile.TarInfo(path)
                tarinfo.size = fileobj.size
                mtime = self._artifact.created_at.timestamp()
                tarinfo.mtime = mtime
                self._tar.addfile(tarinfo, file)
        elif self._tar_closed:
            # No data read and the tar file is closed: stop iterating
            raise StopIteration
        elif len(self._files) == 0:
            # Once the tar file is finalized, some more data might be read on
            # the next iteration
            self.flush_tar()
        else:
            raise AssertionError  # pragma: no cover

        return b""

    def close(self) -> None:
        """Release all resources."""
        if self._closed:
            return
        if not self._tar_closed:
            self.flush_tar()
        self._tar_file.close()
        shutil.rmtree(self._temp_directory)
        self._closed = True
