#  Copyright (c) ZenML GmbH 2023. 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.
"""Pipeline build utilities."""

import hashlib
import platform
import time
from typing import (
    TYPE_CHECKING,
    Dict,
    List,
    Optional,
    Union,
)
from uuid import UUID

import zenml
from zenml.client import Client
from zenml.code_repositories import BaseCodeRepository
from zenml.logger import get_logger
from zenml.models import (
    BuildItem,
    PipelineBuildBase,
    PipelineBuildRequest,
    PipelineBuildResponse,
    PipelineSnapshotBase,
    StackResponse,
)
from zenml.stack import Stack
from zenml.utils import docker_utils, source_utils
from zenml.utils.pipeline_docker_image_builder import (
    PipelineDockerImageBuilder,
)

if TYPE_CHECKING:
    from zenml.code_repositories import LocalRepositoryContext
    from zenml.config.build_configuration import BuildConfiguration
    from zenml.config.docker_settings import DockerSettings

logger = get_logger(__name__)


def get_docker_settings(
    snapshot: "PipelineSnapshotBase",
) -> List["DockerSettings"]:
    """Get a list of all relevant Docker settings for the snapshot.

    Args:
        snapshot: The snapshot.

    Returns:
        A list of all relevant Docker settings.
    """
    docker_settings = [
        step.config.docker_settings
        for step in snapshot.step_configurations.values()
    ]
    if snapshot.is_dynamic:
        # In dynamic pipelines, there is always an orchestration container for
        # which a Docker image needs to be built.
        docker_settings.append(snapshot.pipeline_configuration.docker_settings)

    return docker_settings


def build_required(snapshot: "PipelineSnapshotBase") -> bool:
    """Checks whether a build is required for the snapshot and active stack.

    Args:
        snapshot: The snapshot for which to check.

    Returns:
        If a build is required.
    """
    stack = Client().active_stack
    return bool(stack.get_docker_builds(snapshot=snapshot))


def requires_included_code(
    snapshot: "PipelineSnapshotBase",
    code_repository: Optional["BaseCodeRepository"] = None,
) -> bool:
    """Checks whether the snapshot requires included code.

    Args:
        snapshot: The snapshot.
        code_repository: If provided, this code repository can be used to
            download the code inside the container images.

    Returns:
        If the snapshot requires code included in the container images.
    """
    for docker_settings in get_docker_settings(snapshot=snapshot):
        if docker_settings.local_project_install_command:
            # When installing a local package, we need to include the code
            # files in the container image.
            return True

        if docker_settings.allow_download_from_artifact_store:
            return False

        if docker_settings.allow_download_from_code_repository:
            if code_repository:
                continue

        if docker_settings.allow_including_files_in_images:
            return True

    return False


def requires_download_from_code_repository(
    snapshot: "PipelineSnapshotBase",
) -> bool:
    """Checks whether the snapshot needs to download code from a repository.

    Args:
        snapshot: The snapshot.

    Returns:
        If the snapshot needs to download code from a code repository.
    """
    for docker_settings in get_docker_settings(snapshot=snapshot):
        if docker_settings.allow_download_from_artifact_store:
            return False

        if docker_settings.allow_including_files_in_images:
            return False

        if docker_settings.allow_download_from_code_repository:
            # The other two options are false, which means download from a
            # code repo is required.
            return True

    return False


def code_download_possible(
    snapshot: "PipelineSnapshotBase",
    code_repository: Optional["BaseCodeRepository"] = None,
) -> bool:
    """Checks whether code download is possible for the snapshot.

    Args:
        snapshot: The snapshot.
        code_repository: If provided, this code repository can be used to
            download the code inside the container images.

    Returns:
        Whether code download is possible for the snapshot.
    """
    for docker_settings in get_docker_settings(snapshot=snapshot):
        if docker_settings.local_project_install_command:
            return False

        if docker_settings.allow_download_from_artifact_store:
            continue

        if (
            docker_settings.allow_download_from_code_repository
            and code_repository
        ):
            continue

        return False

    return True


def should_prevent_build_reuse(snapshot: "PipelineSnapshotBase") -> bool:
    """Whether the snapshot prevents a build reuse.

    Args:
        snapshot: The snapshot.

    Returns:
        Whether the snapshot prevents a build reuse.
    """
    return any(
        docker_settings.prevent_build_reuse
        for docker_settings in get_docker_settings(snapshot=snapshot)
    )


def reuse_or_create_pipeline_build(
    snapshot: "PipelineSnapshotBase",
    allow_build_reuse: bool,
    pipeline_id: Optional[UUID] = None,
    build: Union["UUID", "PipelineBuildBase", None] = None,
    code_repository: Optional["BaseCodeRepository"] = None,
) -> Optional["PipelineBuildResponse"]:
    """Loads or creates a pipeline build.

    Args:
        snapshot: The pipeline snapshot for which to load or create the
            build.
        allow_build_reuse: If True, the build is allowed to reuse an
            existing build.
        pipeline_id: Optional ID of the pipeline to reference in the build.
        build: Optional existing build. If given, the build will be fetched
            (or registered) in the database. If not given, a new build will
            be created.
        code_repository: If provided, this code repository can be used to
            download code inside the container images.

    Returns:
        The build response.
    """
    if not build:
        if (
            allow_build_reuse
            and not should_prevent_build_reuse(snapshot=snapshot)
            and not requires_included_code(
                snapshot=snapshot, code_repository=code_repository
            )
            and build_required(snapshot=snapshot)
        ):
            existing_build = find_existing_build(
                snapshot=snapshot, code_repository=code_repository
            )

            if existing_build:
                logger.info(
                    "Reusing existing build `%s` for stack `%s`.",
                    existing_build.id,
                    Client().active_stack.name,
                )
                return existing_build
            else:
                logger.info(
                    "Unable to find a build to reuse. A previous build can be "
                    "reused when the following conditions are met:\n"
                    "  * The existing build was created for the same stack, "
                    "ZenML version and Python version\n"
                    "  * The stack contains a container registry\n"
                    "  * The Docker settings of the pipeline and all its steps "
                    "are the same as for the existing build."
                )

        return create_pipeline_build(
            snapshot=snapshot,
            pipeline_id=pipeline_id,
            code_repository=code_repository,
        )

    if isinstance(build, UUID):
        build_model = Client().zen_store.get_build(build_id=build)
    else:
        build_request = PipelineBuildRequest(
            project=Client().active_project.id,
            stack=Client().active_stack_model.id,
            pipeline=pipeline_id,
            **build.model_dump(),
        )
        build_model = Client().zen_store.create_build(build=build_request)

    verify_custom_build(
        build=build_model,
        snapshot=snapshot,
        code_repository=code_repository,
    )

    return build_model


def find_existing_build(
    snapshot: "PipelineSnapshotBase",
    code_repository: Optional["BaseCodeRepository"] = None,
) -> Optional["PipelineBuildResponse"]:
    """Find an existing build for a snapshot.

    Args:
        snapshot: The snapshot for which to find an existing build.
        code_repository: The code repository that will be used to download
            files in the images.

    Returns:
        The existing build to reuse if found.
    """
    client = Client()
    stack = client.active_stack

    if not stack.container_registry:
        # There can be no non-local builds that we can reuse if there is no
        # container registry in the stack.
        return None

    python_version_prefix = ".".join(platform.python_version_tuple()[:2])
    required_builds = stack.get_docker_builds(snapshot=snapshot)

    if not required_builds:
        return None

    build_checksum = compute_build_checksum(
        required_builds, stack=stack, code_repository=code_repository
    )

    matches = client.list_builds(
        sort_by="desc:created",
        size=1,
        stack_id=stack.id,
        # Until we implement stack versioning, users can still update their
        # stack to update/remove the container registry. In that case, we might
        # try to pull an image from a container registry that we don't have
        # access to. This is why we add an additional check for the container
        # registry ID here. (This is still not perfect as users can update the
        # container registry URI or config, but the best we can do)
        container_registry_id=stack.container_registry.id,
        # The build is local and it's not clear whether the images
        # exist on the current machine or if they've been overwritten.
        # TODO: Should we support this by storing the unique Docker ID for
        #   the image and checking if an image with that ID exists locally?
        is_local=False,
        # The build contains some code which might be different from the
        # local code the user is expecting to run
        contains_code=False,
        zenml_version=zenml.__version__,
        # Match all patch versions of the same Python major + minor
        python_version=f"startswith:{python_version_prefix}",
        checksum=build_checksum,
    )

    if not matches.items:
        return None

    return matches[0]


def create_pipeline_build(
    snapshot: "PipelineSnapshotBase",
    pipeline_id: Optional[UUID] = None,
    code_repository: Optional["BaseCodeRepository"] = None,
) -> Optional["PipelineBuildResponse"]:
    """Builds images and registers the output in the server.

    Args:
        snapshot: The pipeline snapshot.
        pipeline_id: The ID of the pipeline.
        code_repository: If provided, this code repository will be used to
            download inside the build images.

    Returns:
        The build output.

    Raises:
        RuntimeError: If multiple builds with the same key but different
            settings were specified.
    """
    client = Client()
    stack_model = Client().active_stack_model
    stack = client.active_stack
    required_builds = stack.get_docker_builds(snapshot=snapshot)

    if not required_builds:
        logger.debug("No docker builds required.")
        return None

    logger.info(
        "Building Docker image(s) for pipeline `%s`.",
        snapshot.pipeline_configuration.name,
    )
    start_time = time.time()

    docker_image_builder = PipelineDockerImageBuilder()
    images: Dict[str, BuildItem] = {}
    checksums: Dict[str, str] = {}

    for build_config in required_builds:
        combined_key = PipelineBuildBase.get_image_key(
            component_key=build_config.key, step=build_config.step_name
        )
        checksum = build_config.compute_settings_checksum(
            stack=stack, code_repository=code_repository
        )

        if combined_key in images:
            previous_checksum = images[combined_key].settings_checksum

            if previous_checksum != checksum:
                raise RuntimeError(
                    f"Trying to build image for key `{combined_key}` but "
                    "an image for this key was already built with a "
                    "different configuration. This happens if multiple "
                    "stack components specified Docker builds for the same "
                    "key in the `StackComponent.get_docker_builds(...)` "
                    "method. If you're using custom components, make sure "
                    "to provide unique keys when returning your build "
                    "configurations to avoid this error."
                )
            else:
                continue

        if checksum in checksums:
            item_key = checksums[checksum]
            image_name_or_digest = images[item_key].image
            contains_code = images[item_key].contains_code
            requires_code_download = images[item_key].requires_code_download
            dockerfile = images[item_key].dockerfile
            requirements = images[item_key].requirements
        else:
            if build_config.settings.image_tag:
                tag = build_config.settings.image_tag
            else:
                tag = snapshot.pipeline_configuration.name
                if build_config.step_name:
                    tag += f"-{build_config.step_name}"
                tag += f"-{build_config.key}"
                tag = docker_utils.sanitize_tag(tag)

            include_files = build_config.should_include_files(
                code_repository=code_repository,
            )
            requires_code_download = build_config.should_download_files(
                code_repository=code_repository,
            )
            pass_code_repo = (
                build_config.should_download_files_from_code_repository(
                    code_repository=code_repository
                )
            )

            (
                image_name_or_digest,
                dockerfile,
                requirements,
            ) = docker_image_builder.build_docker_image(
                docker_settings=build_config.settings,
                tag=tag,
                stack=stack,
                include_files=include_files,
                entrypoint=build_config.entrypoint,
                extra_files=build_config.extra_files,
                code_repository=code_repository if pass_code_repo else None,
                extra_requirements_files=build_config.extra_requirements_files,
            )
            contains_code = include_files

        images[combined_key] = BuildItem(
            image=image_name_or_digest,
            dockerfile=dockerfile,
            requirements=requirements,
            settings_checksum=checksum,
            contains_code=contains_code,
            requires_code_download=requires_code_download,
        )
        checksums[checksum] = combined_key

    logger.info("Finished building Docker image(s).")

    duration = round(time.time() - start_time)
    is_local = stack.container_registry is None
    contains_code = any(item.contains_code for item in images.values())
    build_checksum = compute_build_checksum(
        required_builds, stack=stack, code_repository=code_repository
    )
    stack_checksum = compute_stack_checksum(stack=stack_model)
    build_request = PipelineBuildRequest(
        project=client.active_project.id,
        stack=stack_model.id,
        pipeline=pipeline_id,
        is_local=is_local,
        contains_code=contains_code,
        images=images,
        zenml_version=zenml.__version__,
        python_version=platform.python_version(),
        checksum=build_checksum,
        stack_checksum=stack_checksum,
        duration=duration,
    )
    return client.zen_store.create_build(build_request)


def compute_build_checksum(
    items: List["BuildConfiguration"],
    stack: "Stack",
    code_repository: Optional["BaseCodeRepository"] = None,
) -> str:
    """Compute an overall checksum for a pipeline build.

    Args:
        items: Items of the build.
        stack: The stack associated with the build. Will be used to gather
            its requirements.
        code_repository: The code repository that will be used to download
            files inside the build. Will be used for its dependency
            specification.

    Returns:
        The build checksum.
    """
    hash_ = hashlib.md5()  # nosec

    for item in items:
        key = PipelineBuildBase.get_image_key(
            component_key=item.key, step=item.step_name
        )

        settings_checksum = item.compute_settings_checksum(
            stack=stack,
            code_repository=code_repository,
        )

        hash_.update(key.encode())
        hash_.update(settings_checksum.encode())

    return hash_.hexdigest()


def verify_local_repository_context(
    snapshot: "PipelineSnapshotBase",
    local_repo_context: Optional["LocalRepositoryContext"],
) -> Optional[BaseCodeRepository]:
    """Verifies the local repository.

    If the local repository exists and has no local changes, code download
    inside the images is possible.

    Args:
        snapshot: The pipeline snapshot.
        local_repo_context: The local repository active at the source root.

    Raises:
        RuntimeError: If the snapshot requires code download but code download
            is not possible.

    Returns:
        The code repository from which to download files for the runs of the
        snapshot, or None if code download is not possible.
    """
    if build_required(snapshot=snapshot):
        if requires_download_from_code_repository(snapshot=snapshot):
            if not local_repo_context:
                raise RuntimeError(
                    "The `DockerSettings` of the pipeline or one of its "
                    "steps specify that code should be downloaded from a "
                    "code repository, but "
                    "there is no code repository active at your current source "
                    f"root `{source_utils.get_source_root()}`."
                )
            elif local_repo_context.is_dirty:
                raise RuntimeError(
                    "The `DockerSettings` of the pipeline or one of its "
                    "steps specify that code should be downloaded from a "
                    "code repository, but "
                    "the code repository active at your current source root "
                    f"`{source_utils.get_source_root()}` has uncommitted "
                    "changes."
                )
            elif local_repo_context.has_local_changes:
                raise RuntimeError(
                    "The `DockerSettings` of the pipeline or one of its "
                    "steps specify that code should be downloaded from a "
                    "code repository, but "
                    "the code repository active at your current source root "
                    f"`{source_utils.get_source_root()}` has unpushed "
                    "changes."
                )

    code_repository = None
    if local_repo_context and not local_repo_context.has_local_changes:
        code_repository = local_repo_context.code_repository

    return code_repository


def verify_custom_build(
    build: "PipelineBuildResponse",
    snapshot: "PipelineSnapshotBase",
    code_repository: Optional["BaseCodeRepository"] = None,
) -> None:
    """Verify a custom build for a pipeline snapshot.

    Args:
        build: The build to verify.
        snapshot: The snapshot for which to verify the build.
        code_repository: Code repository that will be used to download files
            for the snapshot.

    Raises:
        RuntimeError: If the build can't be used for the snapshot.
    """
    stack = Client().active_stack
    required_builds = stack.get_docker_builds(snapshot=snapshot)

    if build.stack and build.stack.id != stack.id:
        logger.warning(
            "The stack `%s` used for the build `%s` is not the same as the "
            "stack `%s` that the pipeline will run on. This could lead "
            "to issues if the stacks have different build requirements.",
            build.stack.name,
            build.id,
            stack.name,
        )

    if build.contains_code:
        logger.warning(
            "The build you specified for this run contains code and will run "
            "with the step code that was included in the Docker images which "
            "might differ from the local code in your client environment."
        )

    if build.requires_code_download:
        if requires_included_code(
            snapshot=snapshot, code_repository=code_repository
        ):
            raise RuntimeError(
                "The `DockerSettings` of the pipeline or one of its "
                "steps specify that code should be included in the Docker "
                "image, but the build you "
                "specified requires code download. Either update your "
                "`DockerSettings` or specify a different build and try "
                "again."
            )

        if (
            requires_download_from_code_repository(snapshot=snapshot)
            and not code_repository
        ):
            raise RuntimeError(
                "The `DockerSettings` of the pipeline or one of its "
                "steps specify that code should be downloaded from a "
                "code repository but "
                "there is no code repository active at your current source "
                f"root `{source_utils.get_source_root()}`."
            )

        if not code_download_possible(
            snapshot=snapshot, code_repository=code_repository
        ):
            raise RuntimeError(
                "The `DockerSettings` of the pipeline or one of its "
                "steps specify that code can not be downloaded from the "
                "artifact store, but the build you specified requires code "
                "download. Either update your `DockerSettings` or specify a "
                "different build and try again."
            )

    if build.checksum:
        build_checksum = compute_build_checksum(
            required_builds, stack=stack, code_repository=code_repository
        )
        if build_checksum != build.checksum:
            logger.warning(
                "The Docker settings used for the build `%s` are "
                "not the same as currently specified for your pipeline. "
                "This means that the build you specified to run this "
                "pipeline might be outdated and most likely contains "
                "outdated requirements.",
                build.id,
            )
    else:
        # No checksum given for the entire build, we manually check that
        # all the images exist and the setting match
        for build_config in required_builds:
            try:
                image = build.get_image(
                    component_key=build_config.key,
                    step=build_config.step_name,
                )
            except KeyError:
                raise RuntimeError(
                    "The build you specified is missing an image for key: "
                    f"{build_config.key}."
                )

            if build_config.compute_settings_checksum(
                stack=stack, code_repository=code_repository
            ) != build.get_settings_checksum(
                component_key=build_config.key, step=build_config.step_name
            ):
                logger.warning(
                    "The Docker settings used to build the image `%s` are "
                    "not the same as currently specified for your pipeline. "
                    "This means that the build you specified to run this "
                    "pipeline might be outdated and most likely contains "
                    "outdated code or requirements.",
                    image,
                )

    if build.is_local:
        logger.warning(
            "You manually specified a local build to run your pipeline. "
            "This might lead to errors if the images don't exist on "
            "your local machine or the image tags have been "
            "overwritten since the original build happened."
        )


def compute_stack_checksum(stack: StackResponse) -> str:
    """Compute a stack checksum.

    Args:
        stack: The stack for which to compute the checksum.

    Returns:
        The checksum.
    """
    hash_ = hashlib.md5()  # nosec

    # This checksum is used to see if the stack has been updated since a build
    # was created for it. We create this checksum not with specific requirements
    # as these might change with new ZenML releases, but they don't actually
    # invalidate those Docker images.
    required_integrations = sorted(
        {
            component.integration
            for components in stack.components.values()
            for component in components
            if component.integration and component.integration != "built-in"
        }
    )
    for integration in required_integrations:
        hash_.update(integration.encode())

    return hash_.hexdigest()


def should_upload_code(
    snapshot: PipelineSnapshotBase,
    build: Optional[PipelineBuildResponse],
    can_download_from_code_repository: bool,
) -> bool:
    """Checks whether the current code should be uploaded for the snapshot.

    Args:
        snapshot: The snapshot.
        build: The build for the snapshot.
        can_download_from_code_repository: Whether the code can be downloaded
            from a code repository.

    Returns:
        Whether the current code should be uploaded for the snapshot.
    """
    if not build:
        # No build means we don't need to download code into a Docker container
        # for step execution. In other remote orchestrators that don't use
        # Docker containers but instead use e.g. Wheels to run, the code should
        # already be included.
        return False

    if build.contains_code:
        return False

    for docker_settings in get_docker_settings(snapshot=snapshot):
        if (
            can_download_from_code_repository
            and docker_settings.allow_download_from_code_repository
        ):
            # No upload needed for this step
            continue

        if docker_settings.allow_download_from_artifact_store:
            return True

    return False


def allows_download_from_code_repository(
    snapshot: PipelineSnapshotBase,
) -> bool:
    """Checks whether a code repository can be used to download code.

    Args:
        snapshot: The snapshot.

    Returns:
        Whether a code repository can be used to download code.
    """
    for docker_settings in get_docker_settings(snapshot=snapshot):
        if docker_settings.allow_download_from_code_repository:
            return True

    return False


def log_code_repository_usage(
    snapshot: PipelineSnapshotBase,
    local_repo_context: "LocalRepositoryContext",
) -> None:
    """Log what the code repository can (not) be used for given a snapshot.

    Args:
        snapshot: The snapshot.
        local_repo_context: The local repository context.
    """
    if build_required(snapshot) and allows_download_from_code_repository(
        snapshot
    ):
        if local_repo_context.is_dirty:
            logger.warning(
                "Unable to use code repository `%s` to download code or track "
                "the commit hash as there are uncommitted or untracked files.",
                local_repo_context.code_repository.name,
            )
        elif local_repo_context.has_local_changes:
            logger.warning(
                "Unable to use code repository `%s` to download code as there "
                "are unpushed commits.",
                local_repo_context.code_repository.name,
            )
        else:
            logger.info(
                "Using code repository `%s` to download code for this run.",
                local_repo_context.code_repository.name,
            )
    elif local_repo_context.is_dirty:
        logger.warning(
            "Unable to use code repository `%s` to track the commit hash as "
            "there are uncommitted or untracked files.",
            local_repo_context.code_repository.name,
        )
