# Copyright 2022 The Pigweed Authors
#
# 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.
"""WORK IN PROGRESS!

This is intended to be a replacement for the proto codegen in proto.bzl, which
relies on the transitive proto compilation support removed from newer versions
of rules_proto_grpc. However, the version checked in here does not yet support,

1. Proto libraries with dependencies in external repositories.
2. Proto libraries with strip_import_prefix or import_prefix attributes.

In addition, nanopb proto files are not yet generated.

TODO(b/234873954): Close these gaps and start using this implementation.

# Overview of implementation

(If you just want to use pw_proto_library, see its docstring; this section is
intended to orient future maintainers.)

Proto code generation is carried out by the _pw_proto_library,
_pw_raw_rpc_proto_library and _pw_nanopb_rpc_proto_library rules using aspects
(https://docs.bazel.build/versions/main/skylark/aspects.html). A
_pw_proto_library has a single proto_library as a dependency, but that
proto_library may depend on other proto_library targets; as a result, the
generated .pwpb.h file #include's .pwpb.h files generated from the dependency
proto_libraries. The aspect propagates along the proto_library dependency
graph, running the proto compiler on each proto_library in the original
target's transitive dependencies, ensuring that we're not missing any .pwpb.h
files at C++ compile time.

Although we have a separate rule for each protocol compiler plugin
(_pw_proto_library, _pw_raw_rpc_proto_library, _pw_nanopb_rpc_proto_library),
they actually share an implementation (_pw _impl_pw_proto_library) and use
similar aspects, all generated by _proto_compiler_aspect. The only difference
between the rules are captured in the PIGWEED_PLUGIN dictonary and the aspect
instantiations (_pw_proto_compiler_aspect, etc).

"""

load("//pw_build:pigweed.bzl", "pw_cc_library")
load("@rules_proto//proto:defs.bzl", "ProtoInfo")
load("//pw_protobuf_compiler:pw_nanopb_cc_library", "pw_nanopb_cc_library")

def pw_proto_library(
        name = "",
        deps = [],
        nanopb_options = None,
        enabled_targets = None):
    """Generate Pigweed proto C++ code.

    This is the only public symbol in this file: everything else is
    implementation details.

    Args:
      name: The name of the target.
      deps: proto_library targets from which to generate Pigweed C++.
      nanopb_options: path to file containing nanopb options, if any
        (https://jpa.kapsi.fi/nanopb/docs/reference.html#proto-file-options).
      enabled_targets: Specifies which libraries should be generated. Libraries
        will only be generated as needed, but unnecessary outputs may conflict
        with other build rules and thus cause build failures. This filter allows
        manual selection of which libraries should be supported by this build
        target in order to prevent such conflicts. The argument, if provided,
        should be a subset of ["pwpb", "nanopb", "raw_rpc", "nanopb_rpc"]. All
        are enabled by default. Note that "nanopb_rpc" relies on "nanopb".

    Example usage:

      proto_library(
        name = "benchmark_proto",
        srcs = [
          "benchmark.proto",
        ],
      )

      pw_proto_library(
        name = "benchmark_pw_proto",
        deps = [":benchmark_proto"],
      )

      pw_cc_binary(
        name = "proto_user",
        srcs = ["proto_user.cc"],
        deps = [":benchmark_pw_proto.pwpb"],
      )

    The pw_proto_library generates the following targets in this example:

    "benchmark_pw_proto.pwpb": C++ library exposing the "benchmark.pwpb.h" header.
    "benchmark_pw_proto.pwpb_rpc": C++ library exposing the
        "benchmark.rpc.pwpb.h" header.
    "benchmark_pw_proto.raw_rpc": C++ library exposing the "benchmark.raw_rpc.h"
        header.
    "benchmark_pw_proto.nanopb": C++ library exposing the "benchmark.pb.h"
        header.
    "benchmark_pw_proto.nanopb_rpc": C++ library exposing the
        "benchmark.rpc.pb.h" header.
    """

    def is_plugin_enabled(plugin):
        return (enabled_targets == None or plugin in enabled_targets)

    if is_plugin_enabled("nanopb"):
        # Use nanopb to generate the pb.h and pb.c files, and the target
        # exposing them.
        pw_nanopb_cc_library(name + ".nanopb", deps, options = nanopb_options)

    # Use Pigweed proto plugins to generate the remaining files and targets.
    for plugin_name, info in PIGWEED_PLUGIN.items():
        if not is_plugin_enabled(plugin_name):
            continue

        name_pb = name + "_pb." + plugin_name

        plugin_rule = info["compiler"]
        plugin_rule(
            name = name_pb,
            deps = deps,
        )

        # The rpc.pb.h header depends on the generated nanopb or pwpb code.
        if info["include_nanopb_dep"]:
            lib_deps = info["deps"] + [":" + name + ".nanopb"]
        elif info["include_pwpb_dep"]:
            lib_deps = info["deps"] + [":" + name + ".pwpb"]
        else:
            lib_deps = info["deps"]

        pw_cc_library(
            name = name + "." + plugin_name,
            hdrs = [name_pb],
            deps = lib_deps,
            linkstatic = 1,
        )

PwProtoInfo = provider(
    "Returned by PW proto compilation aspect",
    fields = {
        "genfiles": "generated C++ header files",
    },
)

def _get_short_path(source):
    return source.short_path

def _get_path(file):
    return file.path

def _proto_compiler_aspect_impl(target, ctx):
    # List the files we will generate for this proto_library target.
    genfiles = []

    for src in target[ProtoInfo].direct_sources:
        path = src.basename[:-len("proto")] + ctx.attr._extension
        genfiles.append(ctx.actions.declare_file(path, sibling = src))

    args = ctx.actions.args()
    args.add("--plugin=protoc-gen-pwpb={}".format(ctx.executable._protoc_plugin.path))
    args.add("--pwpb_out={}".format(ctx.bin_dir.path))
    args.add_joined(
        "--descriptor_set_in",
        target[ProtoInfo].transitive_descriptor_sets,
        join_with = ctx.host_configuration.host_path_separator,
        map_each = _get_path,
    )

    args.add_all(target[ProtoInfo].direct_sources, map_each = _get_short_path)

    ctx.actions.run(
        inputs = depset(target[ProtoInfo].transitive_sources.to_list(), transitive = [target[ProtoInfo].transitive_descriptor_sets]),
        progress_message = "Generating %s C++ files for %s" % (ctx.attr._extension, ctx.label.name),
        tools = [ctx.executable._protoc_plugin],
        outputs = genfiles,
        executable = ctx.executable._protoc,
        arguments = [args],
    )

    transitive_genfiles = genfiles
    for dep in ctx.rule.attr.deps:
        transitive_genfiles += dep[PwProtoInfo].genfiles
    return [PwProtoInfo(genfiles = transitive_genfiles)]

def _proto_compiler_aspect(extension, protoc_plugin):
    """Returns an aspect that runs the proto compiler.

    The aspect propagates through the deps of proto_library targets, running
    the proto compiler with the specified plugin for each of their source
    files. The proto compiler is assumed to produce one output file per input
    .proto file. That file is placed under bazel-bin at the same path as the
    input file, but with the specified extension (i.e., with _extension =
    .pwpb.h, the aspect converts pw_log/log.proto into
    bazel-bin/pw_log/log.pwpb.h).

    The aspect returns a provider exposing all the File objects generated from
    the dependency graph.
    """
    return aspect(
        attr_aspects = ["deps"],
        attrs = {
            "_extension": attr.string(default = extension),
            "_protoc": attr.label(
                default = Label("@com_google_protobuf//:protoc"),
                executable = True,
                cfg = "exec",
            ),
            "_protoc_plugin": attr.label(
                default = Label(protoc_plugin),
                executable = True,
                cfg = "exec",
            ),
        },
        implementation = _proto_compiler_aspect_impl,
        provides = [PwProtoInfo],
    )

def _impl_pw_proto_library(ctx):
    """Implementation of the proto codegen rule.

    The work of actually generating the code is done by the aspect, so here we
    just gather up all the generated files and return them.
    """

    # Note that we don't distinguish between the files generated from the
    # target, and the files generated from its dependencies. We return all of
    # them together, and in pw_proto_library expose all of them as hdrs.
    # Pigweed's plugins happen to only generate .h files, so this works, but
    # strictly speaking we should expose only the files generated from the
    # target itself in hdrs, and place the headers generated from dependencies
    # in srcs. We don't perform layering_check in Pigweed, so this is not a big
    # deal.
    #
    # TODO(b/234873954): Tidy this up.
    all_genfiles = []
    for dep in ctx.attr.deps:
        for f in dep[PwProtoInfo].genfiles:
            all_genfiles.append(f)

    return [DefaultInfo(files = depset(all_genfiles))]

# Instantiate the aspects and rules for generating code using specific plugins.
_pw_proto_compiler_aspect = _proto_compiler_aspect("pwpb.h", "//pw_protobuf/py:plugin")

_pw_proto_library = rule(
    implementation = _impl_pw_proto_library,
    attrs = {
        "deps": attr.label_list(
            providers = [ProtoInfo],
            aspects = [_pw_proto_compiler_aspect],
        ),
    },
)

_pw_pwpb_rpc_proto_compiler_aspect = _proto_compiler_aspect("rpc.pwpb.h", "//pw_rpc/py:plugin_pwpb")

_pw_pwpb_rpc_proto_library = rule(
    implementation = _impl_pw_proto_library,
    attrs = {
        "deps": attr.label_list(
            providers = [ProtoInfo],
            aspects = [_pw_pwpb_rpc_proto_compiler_aspect],
        ),
    },
)

_pw_raw_rpc_proto_compiler_aspect = _proto_compiler_aspect("raw_rpc.pb.h", "//pw_rpc/py:plugin_raw")

_pw_raw_rpc_proto_library = rule(
    implementation = _impl_pw_proto_library,
    attrs = {
        "deps": attr.label_list(
            providers = [ProtoInfo],
            aspects = [_pw_raw_rpc_proto_compiler_aspect],
        ),
    },
)

_pw_nanopb_rpc_proto_compiler_aspect = _proto_compiler_aspect("rpc.pb.h", "//pw_rpc/py:plugin_nanopb")

_pw_nanopb_rpc_proto_library = rule(
    implementation = _impl_pw_proto_library,
    attrs = {
        "deps": attr.label_list(
            providers = [ProtoInfo],
            aspects = [_pw_nanopb_rpc_proto_compiler_aspect],
        ),
    },
)

PIGWEED_PLUGIN = {
    "pwpb": {
        "compiler": _pw_proto_library,
        "deps": [
            "//pw_assert:facade",
            "//pw_containers:vector",
            "//pw_preprocessor",
            "//pw_protobuf",
            "//pw_result",
            "//pw_span",
            "//pw_status",
            "//pw_string:string",
        ],
        "include_nanopb_dep": False,
        "include_pwpb_dep": False,
    },
    "pwpb_rpc": {
        "compiler": _pw_pwpb_rpc_proto_library,
        "deps": [
            "//pw_protobuf:pw_protobuf",
            "//pw_rpc",
            "//pw_rpc/pwpb:client_api",
            "//pw_rpc/pwpb:server_api",
        ],
        "include_nanopb_dep": False,
        "include_pwpb_dep": True,
    },
    "raw_rpc": {
        "compiler": _pw_raw_rpc_proto_library,
        "deps": [
            "//pw_rpc",
            "//pw_rpc/raw:client_api",
            "//pw_rpc/raw:server_api",
        ],
        "include_nanopb_dep": False,
        "include_pwpb_dep": False,
    },
    "nanopb_rpc": {
        "compiler": _pw_nanopb_rpc_proto_library,
        "deps": [
            "//pw_rpc",
            "//pw_rpc/nanopb:client_api",
            "//pw_rpc/nanopb:server_api",
        ],
        "include_nanopb_dep": True,
        "include_pwpb_dep": False,
    },
}