# Copyright 2017 The Bazel Authors. 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
#
#    http://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.
"""Rule for building a Docker image."""

load(
    ":filetype.bzl",
    deb_filetype = "deb",
    docker_filetype = "docker",
    tar_filetype = "tar",
)
load(
    "//tools/build_defs/hash:hash.bzl",
    _hash_tools = "tools",
    _sha256 = "sha256",
)
load(":label.bzl", _string_to_label = "string_to_label")
load(
    ":layers.bzl",
    _assemble_image = "assemble",
    _get_layers = "get_from_target",
    _incr_load = "incremental_load",
    _layer_tools = "tools",
)
load(":list.bzl", "reverse")
load(
    ":path.bzl",
    "dirname",
    "strip_prefix",
    _canonicalize_path = "canonicalize",
    _join_path = "join",
)
load(":serialize.bzl", _serialize_dict = "dict_to_associative_list")

def _build_layer(ctx):
    """Build the current layer for appending it the base layer."""

    layer = ctx.new_file(ctx.label.name + ".layer")
    build_layer = ctx.executable.build_layer
    args = [
        "--output=" + layer.path,
        "--directory=" + ctx.attr.directory,
        "--mode=" + ctx.attr.mode,
    ]

    if ctx.attr.data_path:
        # If data_prefix is specified, then add files relative to that.
        data_path = _join_path(
            dirname(ctx.outputs.out.short_path),
            _canonicalize_path(ctx.attr.data_path),
        )
        args += [
            "--file=%s=%s" % (f.path, strip_prefix(f.short_path, data_path))
            for f in ctx.files.files
        ]
    else:
        # Otherwise, files are added without a directory prefix at all.
        args += [
            "--file=%s=%s" % (f.path, f.basename)
            for f in ctx.files.files
        ]

    args += ["--tar=" + f.path for f in ctx.files.tars]
    args += ["--deb=" + f.path for f in ctx.files.debs if f.path.endswith(".deb")]
    args += [
        "--link=%s:%s" % (k, ctx.attr.symlinks[k])
        for k in ctx.attr.symlinks
    ]
    arg_file = ctx.new_file(ctx.label.name + ".layer.args")
    ctx.file_action(arg_file, "\n".join(args))

    ctx.action(
        executable = build_layer,
        arguments = ["--flagfile=" + arg_file.path],
        inputs = ctx.files.files + ctx.files.tars + ctx.files.debs + [arg_file],
        outputs = [layer],
        use_default_shell_env = True,
        mnemonic = "DockerLayer",
    )
    return layer

# TODO(mattmoor): In a future change, we should establish the invariant that
# base must expose "docker_layers", possibly by hoisting a "docker_load" rule
# from a tarball "base".
def _get_base_artifact(ctx):
    if ctx.files.base:
        if hasattr(ctx.attr.base, "docker_layers"):
            # The base is the first layer in docker_layers if provided.
            return _get_layers(ctx, ctx.attr.base)[0]["layer"]
        if len(ctx.files.base) != 1:
            fail("base attribute should be a single tar file.")
        return ctx.files.base[0]

def _image_config(ctx, layer_names):
    """Create the configuration for a new docker image."""
    config = ctx.new_file(ctx.label.name + ".config")

    label_file_dict = _string_to_label(
        ctx.files.label_files,
        ctx.attr.label_file_strings,
    )

    labels = dict()
    for l in ctx.attr.labels:
        fname = ctx.attr.labels[l]
        if fname[0] == "@":
            labels[l] = "@" + label_file_dict[fname[1:]].path
        else:
            labels[l] = fname

    args = [
        "--output=%s" % config.path,
        "--entrypoint=%s" % ",".join(ctx.attr.entrypoint),
        "--command=%s" % ",".join(ctx.attr.cmd),
        "--labels=%s" % _serialize_dict(labels),
        "--env=%s" % _serialize_dict(ctx.attr.env),
        "--ports=%s" % ",".join(ctx.attr.ports),
        "--volumes=%s" % ",".join(ctx.attr.volumes),
    ]
    if ctx.attr.user:
        args += ["--user=" + ctx.attr.user]
    if ctx.attr.workdir:
        args += ["--workdir=" + ctx.attr.workdir]

    inputs = layer_names
    args += ["--layer=@" + l.path for l in layer_names]

    if ctx.attr.label_files:
        inputs += ctx.files.label_files

    base = _get_base_artifact(ctx)
    if base:
        args += ["--base=%s" % base.path]
        inputs += [base]

    ctx.action(
        executable = ctx.executable.create_image_config,
        arguments = args,
        inputs = inputs,
        outputs = [config],
        use_default_shell_env = True,
        mnemonic = "ImageConfig",
    )
    return config

def _metadata_action(ctx, layer, name, output):
    """Generate the action to create the JSON metadata for the layer."""
    rewrite_tool = ctx.executable.rewrite_tool

    label_file_dict = _string_to_label(
        ctx.files.label_files,
        ctx.attr.label_file_strings,
    )

    labels = dict()
    for l in ctx.attr.labels:
        fname = ctx.attr.labels[l]
        if fname[0] == "@":
            labels[l] = "@" + label_file_dict[fname[1:]].path
        else:
            labels[l] = fname

    args = [
        "--output=%s" % output.path,
        "--layer=%s" % layer.path,
        "--name=@%s" % name.path,
        "--entrypoint=%s" % ",".join(ctx.attr.entrypoint),
        "--command=%s" % ",".join(ctx.attr.cmd),
        "--labels=%s" % _serialize_dict(labels),
        "--env=%s" % _serialize_dict(ctx.attr.env),
        "--ports=%s" % ",".join(ctx.attr.ports),
        "--volumes=%s" % ",".join(ctx.attr.volumes),
    ]
    if ctx.attr.workdir:
        args += ["--workdir=" + ctx.attr.workdir]
    inputs = [layer, rewrite_tool, name]
    if ctx.attr.label_files:
        inputs += ctx.files.label_files

    # TODO(mattmoor): Does this properly handle naked tarballs?
    base = _get_base_artifact(ctx)
    if base:
        args += ["--base=%s" % base.path]
        inputs += [base]
    if ctx.attr.user:
        args += ["--user=" + ctx.attr.user]

    ctx.action(
        executable = rewrite_tool,
        arguments = args,
        inputs = inputs,
        outputs = [output],
        use_default_shell_env = True,
        mnemonic = "RewriteJSON",
    )

def _metadata(ctx, layer, name):
    """Create the metadata for the new docker image."""
    metadata = ctx.new_file(ctx.label.name + ".metadata")
    _metadata_action(ctx, layer, name, metadata)
    return metadata

def _compute_layer_name(ctx, layer):
    """Compute the layer's name.

    This function synthesize a version of its metadata where in place
    of its final name, we use the SHA256 of the layer blob.

    This makes the name of the layer a function of:
      - Its layer's SHA256
      - Its metadata
      - Its parent's name.
    Assuming the parent's name is derived by this same rigor, then
    a simple induction proves the content addressability.

    Args:
      ctx: Rule context.
      layer: The layer's artifact for which to compute the name.
    Returns:
      The artifact that will contains the name for the layer.
    """
    metadata = ctx.new_file(ctx.label.name + ".metadata-name")
    layer_sha = _sha256(ctx, layer)
    _metadata_action(ctx, layer, layer_sha, metadata)
    return _sha256(ctx, metadata)

def _repository_name(ctx):
    """Compute the repository name for the current rule."""
    if ctx.attr.legacy_repository_naming:
        # Legacy behavior, off by default.
        return _join_path(ctx.attr.repository, ctx.label.package.replace("/", "_"))

    # Newer Docker clients support multi-level names, which are a part of
    # the v2 registry specification.
    return _join_path(ctx.attr.repository, ctx.label.package)

def _create_image(ctx, layers, identifier, config, name, metadata, tags):
    """Create the new image."""
    args = [
        "--output=" + ctx.outputs.layer.path,
        "--id=@" + identifier.path,
        "--config=" + config.path,
    ] + ["--tag=" + tag for tag in tags]

    args += ["--layer=@%s=%s" % (l["name"].path, l["layer"].path) for l in layers]
    inputs = [identifier, config] + [l["name"] for l in layers] + [l["layer"] for l in layers]

    if name:
        args += ["--legacy_id=@" + name.path]
        inputs += [name]

    if metadata:
        args += ["--metadata=" + metadata.path]
        inputs += [metadata]

    # If we have been provided a base image, add it.
    if ctx.attr.base and not hasattr(ctx.attr.base, "docker_layers"):
        legacy_base = _get_base_artifact(ctx)
        if legacy_base:
            args += ["--legacy_base=%s" % legacy_base.path]
            inputs += [legacy_base]

    # TODO(mattmoor): Does this properly handle naked tarballs? (excl. above)
    base = _get_base_artifact(ctx)
    if base:
        args += ["--base=%s" % base.path]
        inputs += [base]
    ctx.action(
        executable = ctx.executable.create_image,
        arguments = args,
        inputs = inputs,
        outputs = [ctx.outputs.layer],
        mnemonic = "CreateImage",
    )

def _docker_build_impl(ctx):
    """Implementation for the docker_build rule."""
    layer = _build_layer(ctx)
    layer_sha = _sha256(ctx, layer)

    config = _image_config(ctx, [layer_sha])
    identifier = _sha256(ctx, config)

    name = _compute_layer_name(ctx, layer)
    metadata = _metadata(ctx, layer, name)

    # Construct a temporary name based on the build target.
    tags = [_repository_name(ctx) + ":" + ctx.label.name]

    # creating a partial image so only pass the layers that belong to it
    image_layer = {"layer": layer, "name": layer_sha}
    _create_image(ctx, [image_layer], identifier, config, name, metadata, tags)

    # Compute the layers transitive provider.
    # This must includes all layers of the image, including:
    #  - The layer introduced by this rule.
    #  - The layers transitively introduced by docker_build deps.
    #  - Layers introduced by a static tarball base.
    # This is because downstream tooling should just be able to depend on
    # the availability and completeness of this field.
    layers = [
        {"layer": ctx.outputs.layer, "id": identifier, "name": name},
    ] + _get_layers(ctx, ctx.attr.base)

    # Generate the incremental load statement
    _incr_load(
        ctx,
        layers,
        {
            tag_name: {"name": name, "id": identifier}
            for tag_name in tags
        },
        ctx.outputs.executable,
    )

    _assemble_image(
        ctx,
        reverse(layers),
        {tag_name: name for tag_name in tags},
        ctx.outputs.out,
    )
    runfiles = ctx.runfiles(
        files = ([l["name"] for l in layers] +
                 [l["id"] for l in layers] +
                 [l["layer"] for l in layers]),
    )
    return struct(
        runfiles = runfiles,
        files = depset([ctx.outputs.layer]),
        docker_layers = layers,
    )

docker_build_ = rule(
    implementation = _docker_build_impl,
    attrs = dict({
        "base": attr.label(allow_files = docker_filetype),
        "data_path": attr.string(),
        "directory": attr.string(default = "/"),
        "tars": attr.label_list(allow_files = tar_filetype),
        "debs": attr.label_list(allow_files = deb_filetype),
        "files": attr.label_list(allow_files = True),
        "legacy_repository_naming": attr.bool(default = False),
        "mode": attr.string(default = "0555"),
        "symlinks": attr.string_dict(),
        "entrypoint": attr.string_list(),
        "cmd": attr.string_list(),
        "user": attr.string(),
        "env": attr.string_dict(),
        "labels": attr.string_dict(),
        "ports": attr.string_list(),  # Skylark doesn't support int_list...
        "volumes": attr.string_list(),
        "workdir": attr.string(),
        "repository": attr.string(default = "bazel"),
        # Implicit dependencies.
        "label_files": attr.label_list(
            allow_files = True,
        ),
        "label_file_strings": attr.string_list(),
        "build_layer": attr.label(
            default = Label("//tools/build_defs/pkg:build_tar"),
            cfg = "host",
            executable = True,
            allow_files = True,
        ),
        "create_image": attr.label(
            default = Label("//tools/build_defs/docker:create_image"),
            cfg = "host",
            executable = True,
            allow_files = True,
        ),
        "rewrite_tool": attr.label(
            default = Label("//tools/build_defs/docker:rewrite_json"),
            cfg = "host",
            executable = True,
            allow_files = True,
        ),
        "create_image_config": attr.label(
            default = Label("//tools/build_defs/docker:create_image_config"),
            cfg = "host",
            executable = True,
            allow_files = True,
        ),
    }.items() + _hash_tools.items() + _layer_tools.items()),
    outputs = {
        "out": "%{name}.tar",
        "layer": "%{name}-layer.tar",
    },
    executable = True,
)

# This validates the two forms of value accepted by
# ENTRYPOINT and CMD, turning them into a canonical
# python list form.
#
# The Dockerfile construct:
#   ENTRYPOINT "/foo"
# Results in:
#   "Entrypoint": [
#       "/bin/sh",
#       "-c",
#       "\"/foo\""
#   ],
# Whereas:
#   ENTRYPOINT ["/foo", "a"]
# Results in:
#   "Entrypoint": [
#       "/foo",
#       "a"
#   ],
# NOTE: prefacing a command with 'exec' just ends up with the former
def _validate_command(name, argument):
    if type(argument) == "string":
        return ["/bin/sh", "-c", argument]
    elif type(argument) == "list":
        return argument
    elif argument:
        fail("The %s attribute must be a string or list, if specified." % name)
    else:
        return None

# Produces a new docker image tarball compatible with 'docker load', which
# is a single additional layer atop 'base'.  The goal is to have relatively
# complete support for building docker image, from the Dockerfile spec.
#
# For more information see the 'Config' section of the image specification:
# https://github.com/opencontainers/image-spec/blob/v0.2.0/serialization.md
#
# Only 'name' is required. All other fields have sane defaults.
#
#   docker_build(
#      name="...",
#      visibility="...",
#
#      # The base layers on top of which to overlay this layer,
#      # equivalent to FROM.
#      base="//another/build:rule",
#
#      # The base directory of the files, defaulted to
#      # the package of the input.
#      # All files structure relatively to that path will be preserved.
#      # A leading '/' mean the workspace root and this path is relative
#      # to the current package by default.
#      data_path="...",
#
#      # The directory in which to expand the specified files,
#      # defaulting to '/'.
#      # Only makes sense accompanying one of files/tars/debs.
#      directory="...",
#
#      # The set of archives to expand, or packages to install
#      # within the chroot of this layer
#      files=[...],
#      tars=[...],
#      debs=[...],
#
#      # The set of symlinks to create within a given layer.
#      symlinks = {
#          "/path/to/link": "/path/to/target",
#          ...
#      },
#
#      # https://docs.docker.com/reference/builder/#entrypoint
#      entrypoint="...", or
#      entrypoint=[...],            -- exec form
#
#      # https://docs.docker.com/reference/builder/#cmd
#      cmd="...", or
#      cmd=[...],                   -- exec form
#
#      # https://docs.docker.com/reference/builder/#expose
#      ports=[...],
#
#      # https://docs.docker.com/reference/builder/#user
#      # NOTE: the normal directive affects subsequent RUN, CMD,
#      # and ENTRYPOINT
#      user="...",
#
#      # https://docs.docker.com/reference/builder/#volume
#      volumes=[...],
#
#      # https://docs.docker.com/reference/builder/#workdir
#      # NOTE: the normal directive affects subsequent RUN, CMD,
#      # ENTRYPOINT, ADD, and COPY, but this attribute only affects
#      # the entry point.
#      workdir="...",
#
#      # https://docs.docker.com/reference/builder/#env
#      env = {
#         "var1": "val1",
#         "var2": "val2",
#         ...
#         "varN": "valN",
#      },
#   )
def docker_build(**kwargs):
    """Package a docker image.

    This rule generates a sequence of genrules the last of which is named 'name',
    so the dependency graph works out properly.  The output of this rule is a
    tarball compatible with 'docker save/load' with the structure:
      {layer-name}:
        layer.tar
        VERSION
        json
      {image-config-sha256}.json
      ...
      manifest.json
      repositories
      top     # an implementation detail of our rules, not consumed by Docker.
    This rule appends a single new layer to the tarball of this form provided
    via the 'base' parameter.

    The images produced by this rule are always named 'bazel/tmp:latest' when
    loaded (an internal detail).  The expectation is that the images produced
    by these rules will be uploaded using the 'docker_push' rule below.

    Args:
      **kwargs: See above.
    """
    if "cmd" in kwargs:
        kwargs["cmd"] = _validate_command("cmd", kwargs["cmd"])
    for reserved in ["label_files", "label_file_strings"]:
        if reserved in kwargs:
            fail("reserved for internal use by docker_build macro", attr = reserved)
    if "labels" in kwargs:
        files = sorted(depset([v[1:] for v in kwargs["labels"].values() if v[0] == "@"]))
        kwargs["label_files"] = files
        kwargs["label_file_strings"] = files
    if "entrypoint" in kwargs:
        kwargs["entrypoint"] = _validate_command("entrypoint", kwargs["entrypoint"])
    docker_build_(**kwargs)