# Copyright 2026 Google LLC
#
# 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.
#

"""Starlark build rules for CBOR code generation."""

load("@protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
load("@rules_cc//cc:defs.bzl", "cc_library")

_PROTOC = "@protobuf//:protoc"
_PLUGIN = "//cbor:cbor_generator_main"
_CLANG_FORMAT = "@protobuf//:protoc"

def cc_cbor_proto_library(
        *,
        name,
        srcs,
        deps = [],
        **kwargs):
    """Generates C++ CBOR code from a proto_library.

    Args:
      name: name of the generated cc_library
      srcs: proto_library rules that wrap the *.proto files; must contain exactly one entry
      deps: cc_proto_library rules that wrap the proto_library above
      **kwargs: extra attributes for the generated cc_library
    """
    if len(srcs) != 1:
        fail("srcs must contain exactly one entry", attr = "srcs")

    proto_rule_name = srcs[0]
    if proto_rule_name.startswith(":"):
        proto_rule_name = proto_rule_name[1:]

    proto_rule = native.existing_rule(proto_rule_name)
    if not proto_rule:
        fail("'%s' not found. Define it before cbor_proto_library." % proto_rule_name)

    if proto_rule["kind"] != "proto_library":
        fail("src '%s' has kind '%s', expecting 'proto_library'." % (proto_rule_name, proto_rule["kind"]))

    proto_srcs = list(proto_rule["srcs"])

    generated_srcs = [src.removesuffix(".proto") + ".cbor.cc" for src in proto_srcs]
    generated_hdrs = [src.removesuffix(".proto") + ".cbor.h" for src in proto_srcs]

    # Automatically infer corresponding `_cbor_proto` dependencies for the `proto_library` dependencies
    # within the same package. For example, if the proto_library depends on :foo_proto, we check if
    # :foo_cbor_proto exists in the same package and add it to ensure transitive header resolution.
    auto_deps = []
    current_pkg = "//" + native.package_name() + ":"
    for dep in proto_rule.get("deps", []):
        target_name = None
        if dep.startswith(current_pkg):
            target_name = dep.removeprefix(current_pkg)
        elif dep.startswith(":"):
            target_name = dep[1:]

        if target_name and target_name.endswith("_proto"):
            cbor_name = target_name.removesuffix("_proto") + "_cbor_proto"
            if native.existing_rule(cbor_name):
                auto_deps.append(":" + cbor_name)

    generated_target_name = name + "_gen"

    _generate_cbor_code(
        name = generated_target_name,
        srcs = srcs,
        generated_files = generated_srcs + generated_hdrs,
        visibility = ["//visibility:private"],
    )

    cc_library(
        name = name,
        srcs = generated_srcs,
        hdrs = generated_hdrs,
        deps = deps + auto_deps + [
            "//cbor",
            "@abseil-cpp//absl/status",
            "@abseil-cpp//absl/status:status_macros",
            "@abseil-cpp//absl/status:statusor",
            "@libcppbor",
        ],
        **kwargs
    )

def _generate_cbor_code_impl(ctx):
    proto_sources = depset(
        [f for dep in ctx.attr.srcs for f in dep[ProtoInfo].direct_sources],
    )
    proto_imports = depset(
        transitive = [dep[ProtoInfo].transitive_sources for dep in ctx.attr.srcs],
    )

    unique_roots = {}
    all_files = depset(transitive = [proto_imports, proto_sources]).to_list()
    for f in all_files:
        idx = f.path.rfind("google/protobuf/")
        if idx != -1:
            root = f.path[:idx]
            if root:
                unique_roots[root] = True

    protoc_paths = ["-I."] + ["-I" + r for r in unique_roots] + ["-I" + f.short_path + "=" + f.path for f in all_files]

    args = ctx.actions.args()
    args.add("--plugin=protoc-gen-cbor=" + ctx.executable._plugin.path)
    args.add_all(protoc_paths)
    args.add("--cbor_out", ctx.bin_dir.path)
    args.add_all(proto_sources)
    args.set_param_file_format("multiline")
    args.use_param_file(param_file_arg = "@%s")

    output_paths = [f.path for f in ctx.outputs.generated_files]
    outputs_str = " ".join(output_paths)

    command = "protoc_path=$1; clang_format_path=$2; shift 2; \"$protoc_path\" \"$@\""

    ctx.actions.run_shell(
        arguments = [ctx.executable._proto_compiler.path, ctx.executable._clang_format.path, args],
        command = command,
        inputs = depset(transitive = [proto_sources, proto_imports]),
        mnemonic = "GenerateCborCode",
        outputs = ctx.outputs.generated_files,
        progress_message = "Generating and formatting CBOR code for %{label}",
        tools = depset([
            ctx.executable._proto_compiler,
            ctx.executable._plugin,
            ctx.executable._clang_format,
        ]),
    )
    return [DefaultInfo(files = depset(ctx.outputs.generated_files))]

_generate_cbor_code = rule(
    attrs = {
        "srcs": attr.label_list(
            allow_empty = False,
            mandatory = True,
            providers = [ProtoInfo],
        ),
        "_proto_compiler": attr.label(
            default = Label(_PROTOC),
            executable = True,
            cfg = "exec",
        ),
        "_plugin": attr.label(
            default = Label(_PLUGIN),
            executable = True,
            cfg = "exec",
        ),
        "_clang_format": attr.label(
            default = Label(_CLANG_FORMAT),
            executable = True,
            cfg = "exec",
        ),
        "generated_files": attr.output_list(),
    },
    implementation = _generate_cbor_code_impl,
)