#  Copyright (c) ZenML GmbH 2021. 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.
"""Utility functions relating to Docker."""

import json
import os
import re
from typing import (
    Any,
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
    cast,
)

from docker.client import DockerClient
from docker.errors import DockerException
from docker.utils import build as docker_build_utils

from zenml.io import fileio
from zenml.logger import get_logger
from zenml.utils import io_utils, string_utils

logger = get_logger(__name__)


def check_docker() -> bool:
    """Checks if Docker is installed and running.

    Returns:
        `True` if Docker is installed, `False` otherwise.
    """
    # Try to ping Docker, to see if it's running
    try:
        docker_client = DockerClient.from_env()
        docker_client.ping()
        return True
    except Exception:
        logger.debug("Docker is not running.", exc_info=True)

    return False


def _parse_dockerignore(dockerignore_path: str) -> List[str]:
    """Parses a dockerignore file and returns a list of patterns to ignore.

    Args:
        dockerignore_path: Path to the dockerignore file.

    Returns:
        List of patterns to ignore.
    """
    try:
        file_content = io_utils.read_file_contents_as_string(dockerignore_path)
    except FileNotFoundError:
        logger.warning(
            "Unable to find dockerignore file at path '%s'.", dockerignore_path
        )
        return []

    exclude_patterns = []
    for line in file_content.split("\n"):
        line = line.strip()
        if line and not line.startswith("#"):
            exclude_patterns.append(line)

    return exclude_patterns


def _create_custom_build_context(
    dockerfile_contents: str,
    build_context_root: Optional[str] = None,
    dockerignore: Optional[str] = None,
    extra_files: Sequence[Tuple[str, str]] = (),
) -> Any:
    """Creates a docker build context.

    Args:
        build_context_root: Path to a directory that will be sent to the
            docker daemon as build context.
        dockerfile_contents: File contents of the Dockerfile to use for the
            build.
        dockerignore: Optional path to a dockerignore file. If no value is
            given, the .dockerignore in the root of the build context will be
            used if it exists. Otherwise, all files inside `build_context_root`
            are included in the build context.
        extra_files: Additional files to include in the build context. The
            files should be passed as a tuple
            (filepath_inside_build_context, file_content) and will overwrite
            existing files in the build context if they share the same path.

    Returns:
        Docker build context that can be passed when building a docker image.
    """
    if build_context_root:
        logger.info(
            "Creating Docker build context from directory `%s`.",
            os.path.abspath(build_context_root),
        )
        exclude_patterns = []
        default_dockerignore_path = os.path.join(
            build_context_root, ".dockerignore"
        )
        if dockerignore or fileio.exists(default_dockerignore_path):
            if not dockerignore:
                dockerignore = default_dockerignore_path
            logger.info(
                "Using dockerignore file `%s` to create docker build context.",
                dockerignore,
            )
            exclude_patterns = _parse_dockerignore(dockerignore)
        else:
            logger.info(
                "No `.dockerignore` found, including all files inside build "
                "context.",
            )

        logger.debug(
            "Exclude patterns for creating docker build context: %s",
            exclude_patterns,
        )
        no_ignores_found = not exclude_patterns

        files = docker_build_utils.exclude_paths(
            build_context_root, patterns=exclude_patterns
        )
    else:
        # No build context path given, we create a build context with only
        # the Dockerfile included
        logger.info(
            "No build context root specified, not including any files in the "
            "Docker build context."
        )
        no_ignores_found = False
        files = []

    extra_files = [*extra_files, ("Dockerfile", dockerfile_contents)]
    context = docker_build_utils.create_archive(
        root=build_context_root,
        files=sorted(files),
        gzip=False,
        extra_files=extra_files,
    )

    build_context_size = os.path.getsize(context.name)
    if build_context_size > 50 * 1024 * 1024 and no_ignores_found:
        # The build context exceeds 50MiB and we didn't find any excludes
        # in dockerignore files -> remind to specify a .dockerignore file
        logger.warning(
            "Build context size for docker image: `%s`. If you believe this is "
            "unreasonably large, make sure to include a `.dockerignore` file "
            "at the root of your build context `%s` or specify a custom file "
            "in the Docker configuration when defining your pipeline.",
            string_utils.get_human_readable_filesize(build_context_size),
            default_dockerignore_path,
        )

    return context


def build_image(
    image_name: str,
    dockerfile: Union[str, List[str]],
    build_context_root: Optional[str] = None,
    dockerignore: Optional[str] = None,
    extra_files: Sequence[Tuple[str, str]] = (),
    **custom_build_options: Any,
) -> None:
    """Builds a docker image.

    Args:
        image_name: The name to use for the built docker image.
        dockerfile: Path to a dockerfile or a list of strings representing the
            Dockerfile lines/commands.
        build_context_root: Optional path to a directory that will be sent to
            the Docker daemon as build context. If left empty, the Docker build
            context will be empty.
        dockerignore: Optional path to a dockerignore file. If no value is
            given, the .dockerignore in the root of the build context will be
            used if it exists. Otherwise, all files inside `build_context_root`
            are included in the build context.
        extra_files: Additional files to include in the build context. The
            files should be passed as a tuple
            (filepath_inside_build_context, file_content) and will overwrite
            existing files in the build context if they share the same path.
        **custom_build_options: Additional options that will be passed
            unmodified to the Docker build call when building the image. You
            can use this to for example specify build args or a target stage.
            See https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
            for a full list of available options.
    """
    if isinstance(dockerfile, str):
        dockerfile_contents = io_utils.read_file_contents_as_string(dockerfile)
        logger.info("Using Dockerfile `%s`.", os.path.abspath(dockerfile))
    else:
        dockerfile_contents = "\n".join(dockerfile)

    build_context = _create_custom_build_context(
        dockerfile_contents=dockerfile_contents,
        build_context_root=build_context_root,
        dockerignore=dockerignore,
        extra_files=extra_files,
    )

    build_options = {
        "rm": False,  # don't remove intermediate containers to improve caching
        "pull": True,  # always pull parent images
        **custom_build_options,
    }

    logger.info("Building Docker image `%s`.", image_name)
    logger.debug("Docker build options: %s", build_options)

    logger.info("Building the image might take a while...")

    docker_client = _try_get_docker_client_from_env()

    # We use the client api directly here, so we can stream the logs
    output_stream = docker_client.images.client.api.build(
        fileobj=build_context,
        custom_context=True,
        tag=image_name,
        **build_options,
    )
    _process_stream(output_stream)

    logger.info("Finished building Docker image `%s`.", image_name)


def push_image(
    image_name: str, docker_client: Optional[DockerClient] = None
) -> str:
    """Pushes an image to a container registry.

    Args:
        image_name: The full name (including a tag) of the image to push.
        docker_client: Optional Docker client to use for pushing the image. If
            no client is given, a new client will be created using the default
            Docker environment.

    Returns:
        The Docker repository digest of the pushed image.

    Raises:
        RuntimeError: If fetching the repository digest of the image failed.
    """
    logger.info("Pushing Docker image `%s`.", image_name)
    docker_client = docker_client or _try_get_docker_client_from_env()
    output_stream = docker_client.images.push(image_name, stream=True)
    aux_info = _process_stream(output_stream)
    logger.info("Finished pushing Docker image.")

    image_name_without_tag, _ = image_name.rsplit(":", maxsplit=1)
    prefix_candidates = [f"{image_name_without_tag}@"]

    if image_name_without_tag.startswith(("index.docker.io/", "docker.io/")):
        # When looking for the repo digest later, Docker sometimes removes the
        # index prefix, so we make sure to check for a digest with and without.
        image_name_without_index = image_name_without_tag.split(
            "/", maxsplit=1
        )[1]
        prefix_candidates.append(f"{image_name_without_index}@")

    image = docker_client.images.get(image_name)
    repo_digests: List[str] = image.attrs["RepoDigests"]

    for digest in repo_digests:
        if digest.startswith(tuple(prefix_candidates)):
            return digest

    for info in reversed(aux_info):
        try:
            repo_digest = info["Digest"]
            return f"{image_name_without_tag}@{repo_digest}"
        except KeyError:
            pass
    else:
        raise RuntimeError(
            f"Unable to find repo digest after pushing image {image_name}."
        )


def tag_image(image_name: str, target: str) -> None:
    """Tags an image.

    Args:
        image_name: The name of the image to tag.
        target: The full target name including a tag.
    """
    docker_client = _try_get_docker_client_from_env()
    image = docker_client.images.get(image_name)
    image.tag(target)


def get_image_digest(image_name: str) -> Optional[str]:
    """Gets the digest of an image.

    Args:
        image_name: Name of the image to get the digest for.

    Returns:
        Returns the repo digest for the given image if there exists exactly one.
        If there are zero or multiple repo digests, returns `None`.
    """
    docker_client = _try_get_docker_client_from_env()

    image = docker_client.images.get(image_name)
    repo_digests = image.attrs["RepoDigests"]

    if len(repo_digests) == 1:
        return cast(str, repo_digests[0])
    else:
        logger.debug(
            "Found zero or more repo digests for docker image '%s': %s",
            image_name,
            repo_digests,
        )
        return None


def is_local_image(image_name: str) -> bool:
    """Returns whether an image was pulled from a registry or not.

    Args:
        image_name: Name of the image to check.

    Returns:
        `True` if the image was pulled from a registry, `False` otherwise.
    """
    docker_client = _try_get_docker_client_from_env()
    images = docker_client.images.list(name=image_name)
    if images:
        # An image with this name is available locally -> now check whether it
        # was pulled from a repo or built locally (in which case the repo
        # digest is empty)
        return get_image_digest(image_name) is None
    else:
        # no image with this name found locally
        return False


def _try_get_docker_client_from_env() -> DockerClient:
    """Tries to create a Docker client from the environment.

    Raises:
        RuntimeError: If creating a Docker client from the environment failed.

    Returns:
        A Docker client created from the environment.
    """
    try:
        return DockerClient.from_env()
    except DockerException as e:
        raise RuntimeError(
            "Could not create a Docker client from the environment. Is your Docker daemon running?"
        ) from e


def _process_stream(stream: Iterable[bytes]) -> List[Dict[str, Any]]:
    """Processes the output stream of a docker command call.

    Args:
        stream: The output stream of a docker command call.

    Raises:
        RuntimeError: If there was an error while running the docker command.

    Returns:
        Auxiliary information which was part of the stream.
    """
    auxiliary_info = []

    for element in stream:
        lines = element.decode("utf-8").strip().split("\n")

        for line in lines:
            try:
                line_json = json.loads(line)
                if "error" in line_json:
                    raise RuntimeError(f"Docker error: {line_json['error']}.")
                elif "stream" in line_json:
                    text = line_json["stream"].strip()
                    if "ERROR" in text:
                        logger.error(text)
                    elif re.match(r"^Step [0-9]+/[0-9]+", text):
                        logger.info(text)
                    else:
                        logger.debug(text)
                elif "aux" in line_json:
                    auxiliary_info.append(line_json["aux"])
                else:
                    pass
            except json.JSONDecodeError as error:
                logger.warning(
                    "Failed to decode json for line '%s': %s", line, error
                )

    return auxiliary_info


def sanitize_tag(tag: str) -> str:
    """Sanitize a Docker tag.

    Args:
        tag: The tag to sanitize.

    Raises:
        ValueError: If the tag is empty.

    Returns:
        The sanitized tag.
    """
    tag = re.sub(r"[^a-zA-Z0-9_.\-]", "", tag)
    # Tags can not start with a dot or a dash
    tag = re.sub(r"^[-.]", "_", tag)

    if not tag:
        raise ValueError("Docker image tag cannot be empty.")

    return tag[:128]
