# Copyright 2018-2019 QuantumBlack Visual Analytics Limited
#
# 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
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND
# NONINFRINGEMENT. IN NO EVENT WILL THE LICENSOR OR OTHER CONTRIBUTORS
# BE LIABLE FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF, OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# The QuantumBlack Visual Analytics Limited (“QuantumBlack”) name and logo
# (either separately or in combination, “QuantumBlack Trademarks”) are
# trademarks of QuantumBlack. The License does not grant you any right or
# license to the QuantumBlack Trademarks. You may not use the QuantumBlack
# Trademarks or any confusingly similar mark as a trademark for your product,
#     or use the QuantumBlack Trademarks in any other manner that might cause
# confusion in the marketplace, including but not limited to in advertising,
# on websites, or on software.
#
# See the License for the specific language governing permissions and
# limitations under the License.

""" Utilities for use with click docker commands """

import os
import re
import shutil
import socket
import subprocess
from importlib import import_module
from itertools import chain
from pathlib import Path, PurePosixPath
from subprocess import DEVNULL, PIPE
from typing import List, Sequence, Tuple, Union

from click import ClickException, secho


def check_docker_image_exists(image: str):
    """
    Check that the specified Docker image exists locally.

    Args:
        image: Docker image name.

    Raises:
        ClickException: If specified Docker image was not found.

    """
    command = ["docker", "images", "-q", image]
    res = subprocess.run(command, stdout=PIPE, stderr=DEVNULL)
    if not res.stdout:
        cmd = "kedro docker build --image {0}".format(image)
        raise ClickException(
            "Unable to find image `{0}` locally. Please build it first by running:\n"
            "{1}".format(image, cmd)
        )


def _list_docker_volumes(host_root: str, container_root: str, volumes: Sequence[str]):
    """
    Generate volume mount options as ``docker run`` command expects it.

    Args:
        host_root: Path project root on the host.
        container_root: Path to project root in the container (e.g., `/home/kedro`).
        volumes: List of volumes to mount.

    Yields:
        Tuple[str]: Tuple of the form: ('-v', 'host_path:container_path')

    """
    host_root = Path(host_root).resolve()
    container_root = PurePosixPath(container_root)
    for volume in volumes:
        hpath = host_root / volume  # host path
        cpath = PurePosixPath(container_root) / volume  # container path
        yield "-v", str(hpath) + ":" + str(cpath)


# pylint: disable=too-many-arguments
def compose_docker_run_args(
    host_root: str = None,
    container_root: str = None,
    mount_volumes: Sequence[str] = None,
    required_args: Sequence[Tuple[str, Union[str, None]]] = None,
    optional_args: Sequence[Tuple[str, Union[str, None]]] = None,
    user_args: Sequence[str] = None,
) -> List[str]:
    """
    Make a list of arguments for the docker command.

    Args:
        host_root: Path project root on the host. It must be provided
            if `mount_volumes` are specified, optional otherwise.
        container_root: Path to project root in the container
            (e.g., `/home/kedro/<repo_name>`). It must be
            provided if `mount_volumes` are specified, optional otherwise.
        mount_volumes: List of volumes to be mounted.
        required_args: List of required arguments.
        optional_args: List of optional arguments, these will be added if only
            not present in `user_args` list.
        user_args: List of arguments already specified by the user.
    Raises:
        ClickException: If `mount_volumes` are provided but either `host_root`
            or `container_root` are missing.

    Returns:
        List of arguments for the docker command.
    """

    mount_volumes = mount_volumes or []
    required_args = required_args or []
    optional_args = optional_args or []
    user_args = user_args or []
    split_user_args = {ua.split("=", 1)[0] for ua in user_args}

    def _add_args(name_: str, value_: str = None, force_: bool = False) -> List[str]:
        """
        Add extra args to existing list of CLI args.
        Args:
            name_: Arg name to add.
            value_: Arg value to add, skipped if None.
            force_: Add the argument even if it's present in the current list of args.

        Returns:
            List containing the new args and (optionally) its value or an empty list
                if no values to be added.
        """
        if not force_ and name_ in split_user_args:
            return []
        return [name_] if value_ is None else [name_, value_]

    if mount_volumes:
        if not (host_root and container_root):
            raise ClickException(
                "Both `host_root` and `container_root` must "
                "be specified in `compose_docker_run_args` "
                "call if `mount_volumes` are provided."
            )
        vol_gen = _list_docker_volumes(host_root, container_root, mount_volumes)
        combined_args = list(chain.from_iterable(vol_gen))
    else:
        combined_args = []
    for arg_name, arg_value in required_args:
        combined_args += _add_args(arg_name, arg_value, True)
    for arg_name, arg_value in optional_args:
        combined_args += _add_args(arg_name, arg_value)
    return combined_args + user_args


def make_container_name(image: str, suffix: str = "") -> str:
    """
    Make default container name for the given Docker image.

    Args:
        image: Docker image tag.
        suffix: Suffix to append to the container name.

    Returns:
        Docker container name.
    """
    name = re.sub(r"[^a-zA-Z0-9_.-]+", "-", image)
    if suffix:
        name += "-" + str(suffix)
    return name


def copy_template_files(
    project_path: Path,
    template_path: Path,
    template_files: Sequence[str],
    verbose: bool = False,
):
    """
    If necessary copy files from a template directory into a project directory.

    Args:
        project_path: Destination path.
        template_path: Source path.
        template_files: Files to copy.
        verbose: Echo the names of any created files.

    """
    for file_ in template_files:
        dest = project_path / file_
        if not dest.exists():
            src = template_path / file_
            shutil.copyfile(str(src), str(dest))
            if verbose:
                secho("Creating `{0}`".format(str(dest)))


def get_uid_gid(uid: int = None, gid: int = None) -> Tuple[int, int]:
    """
    Get UID and GID to be passed into the Docker container.
    Defaults to the current user's UID and GID on Unix and (999, 0) on Windows.

    Args:
        uid: Input UID.
        gid: Input GID.

    Returns:
        (UID, GID).
    """

    # Default uid 999 is chosen as the one having potentially the lowest chance
    # of clashing with some existing user in the Docker container.
    _default_uid = 999

    # Default gid 0 corresponds to the root group.
    _default_gid = 0

    if uid is None:
        uid = os.getuid() if os.name == "posix" else _default_uid

    if gid is None:
        gid = (
            import_module("pwd").getpwuid(uid).pw_gid
            if os.name == "posix"
            else _default_gid
        )

    return uid, gid


def add_jupyter_args(run_args: List[str]) -> List[str]:
    """
    Adds `--ip 0.0.0.0` and `--no-browser` options to run args if those are not
    there yet.

    Args:
        run_args: Existing list of run arguments.

    Returns:
        Modified list of run arguments.
    """
    run_args = run_args.copy()
    if not any(arg.split("=", 1)[0] == "--ip" for arg in run_args):
        run_args += ["--ip", "0.0.0.0"]
    if "--no-browser" not in run_args:
        run_args += ["--no-browser"]
    return run_args


def is_port_in_use(port: int) -> bool:
    """Check if the specified TCP port is already in use.

    Args:
        port: Port number.

    Returns:
        True if port is already in use, False otherwise.
    """
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as _s:
        return _s.connect_ex(("0.0.0.0", port)) == 0
