"""Implementation for rules_kustomize"""

load("//hack/build/rules/container:push_info.bzl", "OCIPushInfo")
load("@aspect_bazel_lib//lib:utils.bzl", "propagate_common_rule_attributes")

KustomizeInfo = provider(
    doc = "Contains information about the built Kustomization",
    fields = {
        "deps": "Depset of source files for dependencies of this Kustomization.",
        "image_labels": "A list of pusher references from images",
        "images": "Dict of image references present in the dependencies of this Kustomization keyed by their container_push targets.",
    },
)

def _kustomization_impl(ctx):
    kyaml = ctx.file.kustomization_yaml
    runfiles = ctx.runfiles()

    # Kustomized manifests
    out = ctx.actions.declare_file("{0}_manifests.yaml".format(ctx.label.name))

    # Script that will cat output of the manifests
    printer = ctx.actions.declare_file(ctx.label.name)

    # Instantiate container_push label:image ref dictionary to collect pairings
    # found in the source files of the implementing kustomization and its deps
    image_dict = {}

    image_labels = []

    # Collect images from deps
    for dep in ctx.attr.deps:
        if dep[KustomizeInfo].images:
            dep_images = dep[KustomizeInfo].images
            for image in dep_images:
                if OCIPushInfo in image:
                    image_dict[image] = dep_images[image]
                    image_labels.append(dep_images[image])

    # Collect images from the implementing kustomization
    if ctx.attr.images:
        for image in ctx.attr.images:
            if OCIPushInfo in image:
                image_dict[image] = ctx.attr.images[image]
                image_labels.append(ctx.attr.images[image])

    # Instantiate list of input files for go helper
    input_files = [kyaml] + ctx.files.srcs

    # Instantiate go helper flags and args
    helper_args = ctx.actions.args()
    helper_args.add_joined(["-src-kustomization", ctx.file.kustomization_yaml], join_with = "=")
    helper_args.add_joined(["-out-manifests", out.path], join_with = "=")

    # If any images are found, leverage OCIPushInfo for updated references.
    if len(image_dict) > 0:
        for image in image_dict:
            # Add images entry to helper arguments as a tuple
            existing_ref = image_dict.get(image)
            helper_args.add_joined([existing_ref, image[OCIPushInfo].ref.path], join_with = "=")

            # Include captured output with soruce kustomization.yaml for use in the go helper
            input_files.append(image[OCIPushInfo].ref)

    # Inputs for the go helper
    input_deps = depset(
        direct = input_files,
        transitive = [dep[KustomizeInfo].deps for dep in ctx.attr.deps],
    )

    # Run go helper which modifies kustomization to include spec.images
    ctx.actions.run(
        inputs = input_deps,
        outputs = [out],
        arguments = [helper_args],
        executable = ctx.executable._go_helper,
        progress_message = "Substituting images for %s kustomization..." % ctx.label.name,
    )

    # Generate script to invoke all pushers & write manifests to stdout
    # The script changes the directory to that which contains the kustomized manifests
    ctx.actions.write(
        output = printer,
        content = """#!/usr/bin/env bash
cd ../..
echo \"$(cat {manifests})\"
        """.format(
            manifests = out.basename,
        ),
        is_executable = True,
    )

    return [
        DefaultInfo(
            files = depset([out]),
            executable = printer,
            runfiles = runfiles,
        ),
        KustomizeInfo(
            deps = input_deps,
            images = ctx.attr.images if ctx.attr.images else {},
            image_labels = image_dict,
        ),
    ]

_kustomization = rule(
    implementation = _kustomization_impl,
    attrs = {
        "deps": attr.label_list(
            doc = "Direct dependencies of the Kustomization",
            providers = [KustomizeInfo],
        ),
        "images": attr.label_keyed_string_dict(
            doc = "Image references present in the srcs of this Kustomization paired with their container_push targets",
            providers = [OCIPushInfo],
        ),
        "kustomization_yaml": attr.label(
            doc = "Name of the file containing the Kustomization",
            allow_single_file = True,
            mandatory = True,
        ),
        "srcs": attr.label_list(
            doc = "YAML files included in this Kustomization",
            allow_files = True,
        ),
        "_go_helper": attr.label(
            default = Label("//hack/build/rules/kustomize:kustomize_bin"),
            executable = True,
            cfg = "exec",
        ),
        "_kustomize": attr.label(
            default = Label("//hack/tools:kustomize"),
            allow_single_file = True,
            executable = True,
            cfg = "exec",
        ),
    },
    executable = True,
)

def _kustomization_push_impl(ctx):
    # If any images are found, leverage OCIPushInfo for updated references.
    image_labels = ctx.attr.kustomization[KustomizeInfo].image_labels

    pushers = []

    runfiles = ctx.runfiles()
    for image in image_labels:
        runfiles = runfiles.merge(image[OCIPushInfo].runfiles)
        pushers.append(image[OCIPushInfo].pusher)

    # Join all pusher executables collected
    push_commands = "\n".join([pusher.short_path for pusher in pushers])

    # Script that will invoke any pushers
    pusher = ctx.actions.declare_file(ctx.label.name)

    # Generate script to invoke all pushers & write manifests to stdout
    # The script changes the directory to that which contains the kustomized manifests
    ctx.actions.write(
        output = pusher,
        content = """#!/usr/bin/env bash
{pushers}
""".format(
            pushers = push_commands,
        ),
        is_executable = True,
    )
    return [
        DefaultInfo(
            runfiles = runfiles,
            executable = pusher,
        ),
    ]

_kustomization_push = rule(
    implementation = _kustomization_push_impl,
    attrs = {
        "kustomization": attr.label(
            doc = "YAML files included in this Kustomization",
            providers = [KustomizeInfo],
            allow_files = True,
        ),
    },
    executable = True,
)

def kustomization(
        name,
        kustomization_yaml,
        srcs = None,
        deps = None,
        images = {},
        **kwargs):
    """Creates a kustomization from a source files and other kustomization rules as dependencies.

    The kustomization macro takes inputs and creats a _kustomization rule that generates
    the resolved manifests of the kustomization and any other dependent kustomizations.

    The macro also generates a _kustomization_push rule if any image resolutions are included
    in the _kustomization generated by a macro. The push rule can be used to push any
    images to GAR.

    Args:
        name:                 name of the generated kustomization.
                              Gazelle defaults this to the name of the bazel package
        srcs:                 source files to be used in the kustomization
        deps:                 a list of labels pointing to other _kustomization rules
        images:               a dict of label to GAR url key: vals that can be used to resolve the
                              digest/pusher of a container_push target
        kustomization_yaml:   the kustomization.yaml source file for the kustomization
        **kwargs:             kwargs
    """
    forwarded_args = propagate_common_rule_attributes(kwargs)

    # Build the kustomization
    _kustomization(
        name = name,
        srcs = srcs,
        deps = deps,
        images = images,
        kustomization_yaml = kustomization_yaml,
        **forwarded_args
    )

    if len(images) > 0:
        # Add the pusher
        _kustomization_push(
            name = "{0}_push".format(name),
            kustomization = name,
            **forwarded_args
        )

    return