#  Copyright (c) ZenML GmbH 2022. 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.
"""Implementation of Docker image builds to run ZenML pipelines."""

import itertools
import os
import subprocess
import sys
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
)

import zenml
from zenml.config.docker_settings import (
    DockerBuildConfig,
    DockerBuildOptions,
    DockerSettings,
    PythonEnvironmentExportMethod,
    PythonPackageInstaller,
)
from zenml.constants import (
    ENV_ZENML_CONFIG_PATH,
    ENV_ZENML_ENABLE_REPO_INIT_WARNINGS,
)
from zenml.enums import OperatingSystemType
from zenml.integrations.registry import integration_registry
from zenml.logger import get_logger
from zenml.utils import docker_utils, io_utils, source_utils

if TYPE_CHECKING:
    from zenml.code_repositories import BaseCodeRepository
    from zenml.container_registries import BaseContainerRegistry
    from zenml.image_builders import BuildContext
    from zenml.stack import Stack

logger = get_logger(__name__)

DOCKER_IMAGE_WORKDIR = "/app"
DOCKER_IMAGE_ZENML_CONFIG_PATH = f"{DOCKER_IMAGE_WORKDIR}/.zenconfig"

DEFAULT_DOCKER_PARENT_IMAGE = (
    f"zenmldocker/zenml:{zenml.__version__}-"
    f"py{sys.version_info.major}.{sys.version_info.minor}"
)
DEFAULT_ZENML_DOCKER_REPOSITORY = "zenml"

PIP_DEFAULT_ARGS = {
    "no-cache-dir": None,
    "default-timeout": 60,
}
UV_DEFAULT_ARGS = {"no-cache-dir": None}

DEFAULT_PYPROJECT_EXPORT_COMMANDS = [
    "uv export --format=requirements-txt --directory={directory} --no-hashes --no-emit-project",
    "poetry export --format=requirements.txt --directory={directory} --without-hashes",
]


class PipelineDockerImageBuilder:
    """Builds Docker images to run a ZenML pipeline."""

    def build_docker_image(
        self,
        docker_settings: "DockerSettings",
        tag: str,
        stack: "Stack",
        include_files: bool,
        entrypoint: Optional[str] = None,
        extra_files: Optional[Dict[str, str]] = None,
        code_repository: Optional["BaseCodeRepository"] = None,
        extra_requirements_files: Dict[str, List[str]] = {},
    ) -> Tuple[str, Optional[str], Optional[str]]:
        """Builds (and optionally pushes) a Docker image to run a pipeline.

        Use the image name returned by this method whenever you need to uniquely
        reference the pushed image in order to pull or run it.

        Args:
            docker_settings: The settings for the image build.
            tag: The tag to use for the image.
            stack: The stack on which the pipeline will be deployed.
            include_files: Whether to include files in the build context.
            entrypoint: Entrypoint to use for the final image. If left empty,
                no entrypoint will be included in the image.
            extra_files: Extra files to add to the build context. Keys are the
                path inside the build context, values are either the file
                content or a file path.
            code_repository: The code repository from which files will be
                downloaded.
            extra_requirements_files: Extra requirements to install in the
                Docker image. Each key is the name of a Python requirements file
                to be created and the value is the list of requirements to be
                installed.

        Returns:
            A tuple (image_digest, dockerfile, requirements):
            - The Docker image repo digest or local name, depending on whether
            the image was pushed or is just stored locally.
            - Dockerfile will contain the contents of the Dockerfile used to
            build the image.
            - Requirements is a string with a single pip requirement per line.

        Raises:
            RuntimeError: If the stack does not contain an image builder.
            ValueError: If no Dockerfile and/or custom parent image is
                specified and the Docker configuration doesn't require an
                image build.
            ValueError: If the specified Dockerfile does not exist.
        """
        requirements: Optional[str] = None
        dockerfile: Optional[str] = None

        if docker_settings.skip_build:
            assert (
                docker_settings.parent_image
            )  # checked via validator already

            # Should we tag this here and push it to the container registry of
            # the stack to make sure it's always accessible when running the
            # pipeline?
            return docker_settings.parent_image, dockerfile, requirements

        if docker_settings.dockerfile and not os.path.isfile(
            docker_settings.dockerfile
        ):
            raise ValueError(
                "Dockerfile at path "
                f"{os.path.abspath(docker_settings.dockerfile)} not found."
            )

        stack.validate()
        image_builder = stack.image_builder
        if not image_builder:
            raise RuntimeError(
                "Unable to build Docker images without an image builder in the "
                f"stack `{stack.name}`."
            )

        container_registry = stack.container_registry

        build_context_class = image_builder.build_context_class
        target_image_name = self._get_target_image_name(
            docker_settings=docker_settings,
            tag=tag,
            container_registry=container_registry,
        )

        requires_zenml_build = any(
            [
                docker_settings.requirements,
                docker_settings.pyproject_path,
                docker_settings.required_integrations,
                docker_settings.replicate_local_python_environment,
                docker_settings.install_stack_requirements,
                docker_settings.local_project_install_command,
                docker_settings.apt_packages,
                docker_settings.environment,
                include_files,
                entrypoint,
                extra_files,
                extra_requirements_files,
            ]
        )

        # Fallback to the value defined on the stack component if the
        # pipeline configuration doesn't have a configured value
        parent_image = (
            docker_settings.parent_image or DEFAULT_DOCKER_PARENT_IMAGE
        )

        if docker_settings.dockerfile:
            if parent_image != DEFAULT_DOCKER_PARENT_IMAGE:
                logger.warning(
                    "You've specified both a Dockerfile and a custom parent "
                    "image, ignoring the parent image."
                )

            push = (
                not image_builder.is_building_locally
                or not requires_zenml_build
            )

            if requires_zenml_build:
                # We will build an additional image on top of this one later
                # to include user files and/or install requirements. The image
                # we build now will be used as the parent for the next build.
                repository = docker_settings.target_repository
                if not repository:
                    if container_registry:
                        repository = (
                            container_registry.config.default_repository
                        )

                repository = repository or DEFAULT_ZENML_DOCKER_REPOSITORY
                user_image_tag = docker_utils.sanitize_tag(
                    f"{tag}-intermediate-build"
                )
                user_image_name = f"{repository}:{user_image_tag}"
                if push and container_registry:
                    user_image_name = (
                        f"{container_registry.config.uri}/{user_image_name}"
                    )

                parent_image = user_image_name
            else:
                # The image we'll build from the custom Dockerfile will be
                # used directly, so we tag it with the requested target name.
                user_image_name = target_image_name

            build_config = (
                docker_settings.parent_image_build_config
                or DockerBuildConfig()
            )
            build_context = build_context_class(
                root=docker_settings.build_context_root,
                dockerignore_file=build_config.dockerignore,
            )
            build_context.add_file(
                source=docker_settings.dockerfile, destination="Dockerfile"
            )
            logger.info("Building Docker image `%s`.", user_image_name)
            image_name_or_digest = image_builder.build(
                image_name=user_image_name,
                build_context=build_context,
                docker_build_options=build_config.build_options,
                container_registry=container_registry if push else None,
            )

        elif not requires_zenml_build:
            if parent_image == DEFAULT_DOCKER_PARENT_IMAGE:
                raise ValueError(
                    "Unable to run a ZenML pipeline with the given Docker "
                    "settings: No Dockerfile or custom parent image "
                    "specified and no files will be copied or requirements "
                    "installed."
                )
            else:
                # The parent image will be used directly to run the pipeline and
                # needs to be tagged/pushed
                docker_utils.tag_image(parent_image, target=target_image_name)
                if container_registry:
                    image_name_or_digest = container_registry.push_image(
                        target_image_name
                    )
                else:
                    image_name_or_digest = target_image_name

        if requires_zenml_build:
            logger.info("Building Docker image `%s`.", target_image_name)
            build_config = docker_settings.build_config or DockerBuildConfig()

            # Leave the build context empty if we don't want to include any files
            build_context_root = (
                source_utils.get_source_root() if include_files else None
            )
            dockerignore = (
                build_config.dockerignore or docker_settings.dockerignore
            )
            build_context = build_context_class(
                root=build_context_root,
                dockerignore_file=dockerignore,
            )

            requirements_files = self.gather_requirements_files(
                docker_settings=docker_settings,
                stack=stack,
                code_repository=code_repository,
                extra_requirements_files=extra_requirements_files,
            )

            self._add_requirements_files(
                requirements_files=requirements_files,
                build_context=build_context,
            )
            requirements = (
                "\n".join(
                    file_content for _, file_content, _ in requirements_files
                )
                or None
            )

            apt_packages = docker_settings.apt_packages.copy()
            if docker_settings.install_stack_requirements:
                apt_packages += stack.apt_packages

            # include apt packages from all required integrations
            for integration in docker_settings.required_integrations:
                # get the integration
                integration_cls = integration_registry.integrations[
                    integration
                ]
                apt_packages += integration_cls.APT_PACKAGES

            if apt_packages:
                logger.info(
                    "Including apt packages: %s",
                    ", ".join(f"`{p}`" for p in apt_packages),
                )

            if parent_image == DEFAULT_DOCKER_PARENT_IMAGE:
                # The default parent image is static and doesn't require a pull
                # each time
                pull_parent_image = False
            elif docker_settings.dockerfile and (
                not container_registry or image_builder.is_building_locally
            ):
                # We built a custom parent image and there was either no
                # container registry in the stack to push to or the image was
                # built locally and not pushed, so this is a local image
                pull_parent_image = False
            elif not image_builder.is_building_locally:
                # Remote image builders always need to pull the image
                pull_parent_image = True
            else:
                # If the image is local, we don't need to pull it. Otherwise
                # we play it safe and always pull in case the user pushed a new
                # image for the given name and tag
                pull_parent_image = not docker_utils.is_local_image(
                    parent_image
                )

            build_options = build_config.build_options or DockerBuildOptions()
            build_options.pull = pull_parent_image
            if build_options.rm is None:
                build_options.rm = False

            dockerfile = self._generate_zenml_pipeline_dockerfile(
                parent_image=parent_image,
                docker_settings=docker_settings,
                requirements_files=requirements_files,
                apt_packages=apt_packages,
                entrypoint=entrypoint,
            )
            build_context.add_file(destination="Dockerfile", source=dockerfile)

            if extra_files:
                for destination, source in extra_files.items():
                    build_context.add_file(
                        destination=destination, source=source
                    )

            image_name_or_digest = image_builder.build(
                image_name=target_image_name,
                build_context=build_context,
                docker_build_options=build_options,
                container_registry=container_registry,
            )

        return image_name_or_digest, dockerfile, requirements

    @staticmethod
    def _get_target_image_name(
        docker_settings: "DockerSettings",
        tag: str,
        container_registry: Optional["BaseContainerRegistry"] = None,
    ) -> str:
        """Returns the target image name.

        If a container registry is given, the image name will include the
        registry URI

        Args:
            docker_settings: The settings for the image build.
            tag: The tag to use for the image.
            container_registry: Optional container registry to which this
                image will be pushed.

        Returns:
            The docker image name.
        """
        repository = docker_settings.target_repository
        if not repository:
            if container_registry:
                repository = container_registry.config.default_repository

        repository = repository or DEFAULT_ZENML_DOCKER_REPOSITORY

        target_image_name = f"{repository}:{tag}"
        if container_registry:
            target_image_name = (
                f"{container_registry.config.uri}/{target_image_name}"
            )

        return target_image_name

    @classmethod
    def _add_requirements_files(
        cls,
        requirements_files: List[Tuple[str, str, List[str]]],
        build_context: "BuildContext",
    ) -> None:
        """Adds requirements files to the build context.

        Args:
            requirements_files: List of tuples
                (filename, file_content, pip_options).
            build_context: Build context to add the requirements files to.
        """
        for filename, file_content, _ in requirements_files:
            build_context.add_file(source=file_content, destination=filename)

    @staticmethod
    def gather_requirements_files(
        docker_settings: DockerSettings,
        stack: "Stack",
        code_repository: Optional["BaseCodeRepository"] = None,
        log: bool = True,
        extra_requirements_files: Dict[str, List[str]] = {},
    ) -> List[Tuple[str, str, List[str]]]:
        """Gathers and/or generates pip requirements files.

        This method is called in `PipelineDockerImageBuilder.build_docker_image`
        but it is also called by other parts of the codebase, e.g. the
        `AzureMLStepOperator`, which needs to upload the requirements files to
        AzureML where the step image is then built.

        Args:
            docker_settings: Docker settings that specifies which
                requirements to install.
            stack: The stack on which the pipeline will run.
            code_repository: The code repository from which files will be
                downloaded.
            log: If True, will log the requirements.
            extra_requirements_files: Extra requirements to install in the
                Docker image. Each key is the name of a Python requirements file
                to be created and the value is the list of requirements to be
                installed.

        Raises:
            RuntimeError: If the command to export the local python packages
                failed.
            FileNotFoundError: If the specified requirements file does not
                exist.

        Returns:
            List of tuples (filename, file_content, pip_options) of all
            requirements files.
            The files will be in the following order:
            - Packages installed in the local Python environment
            - Requirements defined by stack integrations
            - Extra requirements files
            - Requirements defined by user integrations
            - Requirements exported from a pyproject.toml
            - User-defined requirements
        """
        implicit_requirements = False
        pyproject_path = docker_settings.pyproject_path
        requirements = docker_settings.requirements

        if not any(
            [
                docker_settings.replicate_local_python_environment,
                requirements,
                pyproject_path,
            ]
        ):
            root = source_utils.get_source_root()
            requirements_path = os.path.join(root, "requirements.txt")
            pyproject_file_path = os.path.join(root, "pyproject.toml")

            if os.path.exists(requirements_path):
                implicit_requirements = True
                requirements = requirements_path
            elif os.path.exists(pyproject_file_path):
                implicit_requirements = True
                pyproject_path = pyproject_file_path

        if (
            implicit_requirements
            and docker_settings.disable_automatic_requirements_detection
        ):
            # TODO: This is only temporary to log a warning notifying users
            # that we will soon switch the default behavior to
            # `disable_automatic_requirements_detection=False`. Remove and
            # adjust the logic once we've made this change.
            if log:
                logger.warning(
                    "Detected `requirements.txt` or `pyproject.toml` files in "
                    "the source root. In future versions of ZenML, these will "
                    "be automatically picked up and installed in Docker images "
                    "by default. To disable this behavior and keep the "
                    "current behavior, set "
                    "`DockerSettings.disable_automatic_requirements_detection` "
                    "to `True`. If you want to enable this behavior right away, "
                    "you can do so by setting "
                    "`DockerSettings.disable_automatic_requirements_detection` "
                    "to `False`."
                )
            implicit_requirements = False
            requirements = None
            pyproject_path = None

        requirements_files: List[Tuple[str, str, List[str]]] = []

        # Generate requirements file for the local environment if configured
        if docker_settings.replicate_local_python_environment:
            if isinstance(
                docker_settings.replicate_local_python_environment,
                PythonEnvironmentExportMethod,
            ):
                command = (
                    docker_settings.replicate_local_python_environment.command
                )
            elif isinstance(
                docker_settings.replicate_local_python_environment, bool
            ):
                command = PythonEnvironmentExportMethod.PIP_FREEZE.command
            else:
                command = " ".join(
                    docker_settings.replicate_local_python_environment
                )

            try:
                local_requirements = subprocess.check_output(
                    command,
                    shell=True,  # nosec
                ).decode()
            except subprocess.CalledProcessError as e:
                raise RuntimeError(
                    "Unable to export local python packages."
                ) from e

            requirements_files.append(
                (".zenml_local_requirements", local_requirements, [])
            )
            if log:
                logger.info(
                    "- Including python packages from local environment"
                )

        if docker_settings.install_stack_requirements:
            stack_requirements = stack.requirements()
            if code_repository:
                stack_requirements.update(code_repository.requirements)

            if stack_requirements:
                stack_requirements_list = sorted(stack_requirements)
                stack_requirements_file = "\n".join(stack_requirements_list)
                requirements_files.append(
                    (
                        ".zenml_stack_integration_requirements",
                        stack_requirements_file,
                        [],
                    )
                )
                if log:
                    logger.info(
                        "- Including stack requirements: %s",
                        ", ".join(f"`{r}`" for r in stack_requirements_list),
                    )

        for filename, extra_requirements in extra_requirements_files.items():
            requirements_list = sorted(extra_requirements)
            requirements_file = "\n".join(requirements_list)
            requirements_files.append((filename, requirements_file, []))
            if log:
                logger.info(
                    "- Including extra requirements from file `%s`: %s",
                    filename,
                    ", ".join(f"`{r}`" for r in requirements_list),
                )

        # Generate requirements file for all required integrations
        integration_requirements = set(
            itertools.chain.from_iterable(
                integration_registry.select_integration_requirements(
                    integration_name=integration,
                    target_os=OperatingSystemType.LINUX,
                )
                for integration in docker_settings.required_integrations
            )
        )

        if integration_requirements:
            integration_requirements_list = sorted(integration_requirements)
            integration_requirements_file = "\n".join(
                integration_requirements_list
            )
            requirements_files.append(
                (
                    ".zenml_integration_requirements",
                    integration_requirements_file,
                    [],
                )
            )
            if log:
                logger.info(
                    "- Including integration requirements: %s",
                    ", ".join(f"`{r}`" for r in integration_requirements_list),
                )

        if pyproject_path:
            path = os.path.abspath(pyproject_path)

            if not os.path.exists(path):
                raise FileNotFoundError(
                    f"Pyproject file {path} does not exist."
                )

            def _run_command(command: str) -> str:
                command = command.format(directory=os.path.dirname(path))

                result = subprocess.run(
                    command,
                    capture_output=True,
                    check=True,
                    shell=True,  # nosec
                )
                return result.stdout.decode()

            if docker_settings.pyproject_export_command:
                command = " ".join(docker_settings.pyproject_export_command)

                try:
                    pyproject_requirements = _run_command(command)
                except subprocess.CalledProcessError as e:
                    raise RuntimeError(
                        "Failed to export pyproject dependencies with "
                        f"command `{command}`: {e.stderr.decode()}"
                    )
            else:
                command_errors = {}

                for command in DEFAULT_PYPROJECT_EXPORT_COMMANDS:
                    try:
                        pyproject_requirements = _run_command(command)
                    except subprocess.CalledProcessError as e:
                        command_errors[command] = e.stderr.decode()
                    else:
                        break
                else:
                    raise RuntimeError(
                        "Failed to export pyproject dependencies with the "
                        f"following commands: {command_errors}. Please specify "
                        "a working command to export your pyproject.toml "
                        "dependencies to a requirements.txt formatted file "
                        "using `DockerSettings.pyproject_export_command`."
                    )

            requirements_files.append(
                (".zenml_pyproject_requirements", pyproject_requirements, [])
            )
            if log:
                logger.info(
                    "- %s python packages from file `%s`",
                    "Implicitly including"
                    if implicit_requirements
                    else "Including",
                    path,
                )

        # Generate/Read requirements file for user-defined requirements
        if isinstance(requirements, str):
            path = os.path.abspath(requirements)
            try:
                user_requirements = io_utils.read_file_contents_as_string(path)
            except FileNotFoundError as e:
                raise FileNotFoundError(
                    f"Requirements file {path} does not exist."
                ) from e
            if log:
                logger.info(
                    "- %s user-defined requirements from file `%s`",
                    "Implicitly including"
                    if implicit_requirements
                    else "Including",
                    path,
                )
        elif isinstance(requirements, List):
            user_requirements = "\n".join(requirements)
            if log:
                logger.info(
                    "- Including user-defined requirements: %s",
                    ", ".join(f"`{r}`" for r in requirements),
                )
        else:
            user_requirements = None

        if user_requirements:
            requirements_files.append(
                (".zenml_user_requirements", user_requirements, [])
            )

        return requirements_files

    @staticmethod
    def _generate_zenml_pipeline_dockerfile(
        parent_image: str,
        docker_settings: DockerSettings,
        requirements_files: Sequence[Tuple[str, str, List[str]]] = (),
        apt_packages: Sequence[str] = (),
        entrypoint: Optional[str] = None,
    ) -> str:
        """Generates a Dockerfile.

        Args:
            parent_image: The image to use as parent for the Dockerfile.
            docker_settings: Docker settings for this image build.
            requirements_files: List of tuples that contain three items:
                - the name of a requirements file,
                - the content of that file,
                - options that should be passed to pip when installing the
                    requirements file.
            apt_packages: APT packages to install.
            entrypoint: The default entrypoint command that gets executed when
                running a container of an image created by this Dockerfile.

        Raises:
            ValueError: If an unsupported python package installer was
                configured.

        Returns:
            The generated Dockerfile.
        """
        lines = [f"FROM {parent_image}", f"WORKDIR {DOCKER_IMAGE_WORKDIR}"]

        if (
            docker_settings.build_config
            and docker_settings.build_config.build_options
        ):
            b_args = docker_settings.build_config.build_options.build_args
            if b_args:
                for arg_key in b_args.keys():
                    lines.append(f"ARG {arg_key}")

        for key, value in docker_settings.environment.items():
            lines.append(f"ENV {key.upper()}='{value}'")

        if apt_packages:
            apt_packages = " ".join(f"'{p}'" for p in apt_packages)

            lines.append(
                "RUN apt-get update && apt-get install -y "
                f"--no-install-recommends {apt_packages}"
            )

        if (
            docker_settings.python_package_installer
            == PythonPackageInstaller.PIP
        ):
            install_command = "pip install"
            default_installer_args: Dict[str, Any] = PIP_DEFAULT_ARGS
        elif (
            docker_settings.python_package_installer
            == PythonPackageInstaller.UV
        ):
            lines.append("RUN pip install uv")
            install_command = "uv pip install"
            default_installer_args = UV_DEFAULT_ARGS
        else:
            raise ValueError("Unsupported python package installer.")

        installer_args = {
            **default_installer_args,
            **docker_settings.python_package_installer_args,
        }
        installer_args_string = " ".join(
            f"--{key}" if value is None else f"--{key}={value}"
            for key, value in installer_args.items()
        )
        for file, _, options in requirements_files:
            lines.append(f"COPY {file} .")
            option_string = " ".join(options)

            lines.append(
                f"RUN {install_command} {installer_args_string}"
                f"{option_string} -r {file}"
            )

        lines.append(f"ENV {ENV_ZENML_ENABLE_REPO_INIT_WARNINGS}=False")
        lines.append(
            f"ENV {ENV_ZENML_CONFIG_PATH}={DOCKER_IMAGE_ZENML_CONFIG_PATH}"
        )

        lines.append("COPY . .")
        lines.append("RUN chmod -R a+rw .")

        if docker_settings.local_project_install_command:
            lines.append(
                f"RUN {docker_settings.local_project_install_command}"
            )

        if docker_settings.user:
            # Change file ownership to specified user
            lines.append(f"RUN chown -R {docker_settings.user} .")
            # Switch back to specified user for subsequent instructions
            lines.append(f"USER {docker_settings.user}")

        for key, value in docker_settings.runtime_environment.items():
            lines.append(f"ENV {key.upper()}='{value}'")

        if entrypoint:
            lines.append(f"ENTRYPOINT {entrypoint}")

        return "\n".join(lines)
