#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
#  or implied. See the License for the specific language governing
#  permissions and limitations under the License.
"""Archivable mixin."""

import io
import tarfile
import zipfile
from abc import ABC, abstractmethod
from pathlib import Path
from typing import IO, Any, Dict, Optional

from zenml.io import fileio
from zenml.utils.enum_utils import StrEnum


class ArchiveType(StrEnum):
    """Archive types supported by the ZenML build context."""

    TAR = "tar"
    TAR_GZ = "tar.gz"
    ZIP = "zip"


class Archivable(ABC):
    """Archivable mixin class."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the object.

        Args:
            *args: Unused args for subclasses.
            **kwargs: Unused keyword args for subclasses.
        """
        self._extra_files: Dict[str, str] = {}

    def add_file(self, source: str, destination: str) -> None:
        """Adds a file to the archive.

        Args:
            source: The source of the file to add. This can either be a path
                or the file content.
            destination: The path inside the archive where the file
                should be added.
        """
        if fileio.exists(source):
            with fileio.open(source) as f:
                self._extra_files[destination] = f.read()
        else:
            self._extra_files[destination] = source

    def add_directory(self, source: str, destination: str) -> None:
        """Adds a directory to the archive.

        Args:
            source: Path to the directory.
            destination: The path inside the build context where the directory
                should be added.

        Raises:
            ValueError: If `source` does not point to a directory.
        """
        if not fileio.isdir(source):
            raise ValueError(
                f"Can't add directory {source} to the build context as it "
                "does not exist or is not a directory."
            )

        for dir, _, files in fileio.walk(source):
            dir_path = Path(fileio.convert_to_str(dir))
            for file_name in files:
                file_name = fileio.convert_to_str(file_name)
                file_source = dir_path / file_name
                file_destination = (
                    Path(destination)
                    / dir_path.relative_to(source)
                    / file_name
                )

                with file_source.open("r") as f:
                    self._extra_files[file_destination.as_posix()] = f.read()

    def write_archive(
        self,
        output_file: IO[bytes],
        archive_type: ArchiveType = ArchiveType.TAR_GZ,
    ) -> None:
        """Writes an archive of the build context to the given file.

        Args:
            output_file: The file to write the archive to.
            archive_type: The type of archive to create.
        """
        files = self.get_files()
        extra_files = self.get_extra_files()
        close_fileobj: Optional[Any] = None
        fileobj: Any = output_file

        if archive_type == ArchiveType.ZIP:
            fileobj = zipfile.ZipFile(output_file, "w", zipfile.ZIP_DEFLATED)
        else:
            if archive_type == ArchiveType.TAR_GZ:
                from gzip import GzipFile

                # We don't use the builtin gzip functionality of the `tarfile`
                # library as that one includes the tar filename and creation
                # timestamp in the archive which causes the hash of the resulting
                # file to be different each time. We use this hash to avoid
                # duplicate uploads, which is why we pass empty values for filename
                # and mtime here.
                close_fileobj = fileobj = GzipFile(
                    filename="", mode="wb", fileobj=output_file, mtime=0.0
                )
            fileobj = tarfile.open(mode="w", fileobj=fileobj)

        try:
            with fileobj as af:
                for archive_path, file_path in files.items():
                    if archive_path in extra_files:
                        continue
                    if archive_type == ArchiveType.ZIP:
                        assert isinstance(af, zipfile.ZipFile)
                        af.write(file_path, arcname=archive_path)
                    else:
                        assert isinstance(af, tarfile.TarFile)
                        if info := af.gettarinfo(
                            file_path, arcname=archive_path
                        ):
                            if info.isfile():
                                with open(file_path, "rb") as f:
                                    af.addfile(info, f)
                            else:
                                af.addfile(info, None)

                for archive_path, contents in extra_files.items():
                    contents_encoded = contents.encode("utf-8")

                    if archive_type == ArchiveType.ZIP:
                        assert isinstance(af, zipfile.ZipFile)
                        af.writestr(archive_path, contents_encoded)
                    else:
                        assert isinstance(af, tarfile.TarFile)
                        info = tarfile.TarInfo(archive_path)
                        info.size = len(contents_encoded)
                        af.addfile(info, io.BytesIO(contents_encoded))
        finally:
            if close_fileobj:
                close_fileobj.close()

        output_file.seek(0)

    @abstractmethod
    def get_files(self) -> Dict[str, str]:
        """Gets all regular files that should be included in the archive.

        Returns:
            A dict {path_in_archive: path_on_filesystem} for all regular files
            in the archive.
        """

    def get_extra_files(self) -> Dict[str, str]:
        """Gets all extra files that should be included in the archive.

        Returns:
            A dict {path_in_archive: file_content} for all extra files in the
            archive.
        """
        return self._extra_files.copy()
