#  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.
"""Code utilities."""

import hashlib
import os
import shutil
import sys
import tempfile
from pathlib import Path
from typing import IO, TYPE_CHECKING, Dict, Optional

from zenml.client import Client
from zenml.io import fileio
from zenml.logger import get_logger
from zenml.utils import io_utils, source_utils, string_utils
from zenml.utils.archivable import Archivable, ArchiveType

if TYPE_CHECKING:
    from git.repo.base import Repo

    from zenml.artifact_stores import BaseArtifactStore


logger = get_logger(__name__)


class CodeArchive(Archivable):
    """Code archive class.

    This class is used to archive user code before uploading it to the artifact
    store. If the user code is stored in a Git repository, only files not
    excluded by gitignores will be included in the archive.
    """

    def __init__(self, root: Optional[str] = None) -> None:
        """Initialize the object.

        Args:
            root: Root directory of the archive.
        """
        super().__init__()
        self._root = root

    @property
    def git_repo(self) -> Optional["Repo"]:
        """Git repository active at the code archive root.

        Returns:
            The git repository if available.
        """
        try:
            # These imports fail when git is not installed on the machine
            from git.exc import InvalidGitRepositoryError
            from git.repo.base import Repo
        except ImportError:
            return None

        try:
            git_repo = Repo(path=self._root, search_parent_directories=True)
        except InvalidGitRepositoryError:
            return None

        return git_repo

    def _get_all_files(self, archive_root: str) -> Dict[str, str]:
        """Get all files inside the archive root.

        Args:
            archive_root: The root directory from which to get all files.

        Returns:
            All files inside the archive root.
        """
        all_files = {}
        for root, _, files in os.walk(archive_root):
            for file in files:
                file_path = os.path.join(root, file)
                path_in_archive = os.path.relpath(file_path, archive_root)
                all_files[path_in_archive] = file_path

        return all_files

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

        Raises:
            RuntimeError: If the code archive would not include any files.

        Returns:
            A dict {path_in_archive: path_on_filesystem} for all regular files
            in the archive.
        """
        if not self._root:
            return {}

        all_files = {}

        if repo := self.git_repo:
            try:
                result = repo.git.ls_files(
                    "--cached",
                    "--others",
                    "--modified",
                    "--exclude-standard",
                    self._root,
                )
            except Exception as e:
                logger.warning(
                    "Failed to get non-ignored files from git: %s", str(e)
                )
                all_files = self._get_all_files(archive_root=self._root)
            else:
                for file in result.split():
                    file_path = os.path.join(repo.working_dir, file)
                    path_in_archive = os.path.relpath(file_path, self._root)

                    if os.path.exists(file_path):
                        all_files[path_in_archive] = file_path
        else:
            all_files = self._get_all_files(archive_root=self._root)

        if not all_files:
            raise RuntimeError(
                "The code archive to be uploaded does not contain any files. "
                "This is probably because all files in your source root "
                f"`{self._root}` are ignored by a .gitignore file."
            )

        # Explicitly remove .zen directories as we write an updated version
        # to disk everytime ZenML is called. This updates the mtime of the
        # file, which invalidates the code upload caching. The values in
        # the .zen directory are not needed anyway as we set them as
        # environment variables.
        all_files = {
            path_in_archive: file_path
            for path_in_archive, file_path in sorted(all_files.items())
            if ".zen" not in Path(path_in_archive).parts[:-1]
        }

        return all_files

    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.
        """
        super().write_archive(
            output_file=output_file, archive_type=archive_type
        )
        archive_size = os.path.getsize(output_file.name)
        if archive_size > 20 * 1024 * 1024:
            logger.warning(
                "Code archive size: `%s`. If you believe this is "
                "unreasonably large, make sure to version your code in git and "
                "ignore unnecessary files using a `.gitignore` file.",
                string_utils.get_human_readable_filesize(archive_size),
            )


def compute_file_hash(file: IO[bytes]) -> str:
    """Compute a hash of the content of a file.

    This function will not seek the file before or after the hash computation.
    This means that the content will be computed based on the current cursor
    until the end of the file.

    Args:
        file: The file for which to compute the hash.

    Returns:
        A hash of the file content.
    """
    hash_ = hashlib.sha1()  # nosec

    while True:
        data = file.read(64 * 1024)
        if not data:
            break
        hash_.update(data)

    return hash_.hexdigest()


def upload_code_if_necessary(code_archive: CodeArchive) -> str:
    """Upload code to the artifact store if necessary.

    This function computes a hash of the code to be uploaded, and if an archive
    with the same hash already exists it will not re-upload but instead return
    the path to the existing archive.

    Args:
        code_archive: The code archive to upload.

    Returns:
        The path where the archived code is uploaded.
    """
    artifact_store = Client().active_stack.artifact_store

    with tempfile.NamedTemporaryFile(
        mode="w+b", delete=False, suffix=".tar.gz"
    ) as f:
        code_archive.write_archive(f)
        archive_path = f.name
        archive_hash = compute_file_hash(f)

    upload_dir = os.path.join(artifact_store.path, "code_uploads")
    fileio.makedirs(upload_dir)
    upload_path = os.path.join(upload_dir, f"{archive_hash}.tar.gz")
    upload_path = io_utils.sanitize_remote_path(upload_path)

    if not fileio.exists(upload_path):
        archive_size = string_utils.get_human_readable_filesize(
            os.path.getsize(archive_path)
        )
        logger.info(
            "Uploading code to `%s` (Size: %s).", upload_path, archive_size
        )
        fileio.copy(archive_path, upload_path)
        logger.info("Code upload finished.")
    else:
        logger.info("Code already exists in artifact store, skipping upload.")

    if os.path.exists(archive_path):
        os.remove(archive_path)

    return upload_path


def download_and_extract_code(code_path: str, extract_dir: str) -> None:
    """Download and extract code.

    Args:
        code_path: Path where the code is uploaded.
        extract_dir: Directory where to code should be extracted to.
    """
    download_path = os.path.basename(code_path)
    fileio.copy(code_path, download_path)

    shutil.unpack_archive(filename=download_path, extract_dir=extract_dir)
    os.remove(download_path)


def download_code_from_artifact_store(
    code_path: str, artifact_store: "BaseArtifactStore"
) -> None:
    """Download code from the artifact store.

    Args:
        code_path: Path where the code is stored.
        artifact_store: The artifact store to use for the download.

    Raises:
        RuntimeError: If the code is stored in an artifact store which is
            not active.
    """
    logger.info("Downloading code from artifact store path `%s`.", code_path)

    if not code_path.startswith(artifact_store.path):
        raise RuntimeError(
            "The code is not stored in the artifact store "
            f"{artifact_store.name} that was passed to download it."
        )

    # Make sure we register the artifact store filesystem here so the
    # fileio.copy call will pick up the right credentials
    artifact_store._register()

    extract_dir = os.path.abspath("code")
    os.makedirs(extract_dir, exist_ok=True)

    download_and_extract_code(code_path=code_path, extract_dir=extract_dir)

    source_utils.set_custom_source_root(extract_dir)
    sys.path.insert(0, extract_dir)
    os.chdir(extract_dir)


def _get_notebook_upload_dir(artifact_store: "BaseArtifactStore") -> str:
    """Get the upload directory for code extracted from notebook cells.

    Args:
        artifact_store: The artifact store in which the directory should be.

    Returns:
        The upload directory for code extracted from notebook cells.
    """
    return os.path.join(artifact_store.path, "notebook_code")


def upload_notebook_code(
    artifact_store: "BaseArtifactStore", cell_code: str, file_name: str
) -> None:
    """Upload code extracted from a notebook cell.

    Args:
        artifact_store: The artifact store in which to upload the code.
        cell_code: The notebook cell code.
        file_name: The filename to use for storing the cell code.
    """
    upload_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
    fileio.makedirs(upload_dir)
    upload_path = os.path.join(upload_dir, file_name)

    if not fileio.exists(upload_path):
        with fileio.open(upload_path, "wb") as f:
            f.write(cell_code.encode())

        logger.info("Uploaded notebook cell code to %s.", upload_path)


def download_notebook_code(
    artifact_store: "BaseArtifactStore", file_name: str, download_path: str
) -> None:
    """Download code extracted from a notebook cell.

    Args:
        artifact_store: The artifact store from which to download the code.
        file_name: The name of the code file.
        download_path: The local path where the file should be downloaded to.

    Raises:
        FileNotFoundError: If no file with the given filename exists in this
            artifact store.
    """
    code_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
    code_path = os.path.join(code_dir, file_name)

    if not fileio.exists(code_path):
        raise FileNotFoundError(
            f"Notebook code at path {code_path} not found."
        )

    fileio.copy(code_path, download_path)
