"""Generates and compiles Python gRPC stubs from proto_library rules."""

load("@rules_proto//proto:defs.bzl", "ProtoInfo")
load(
    "//bazel:protobuf.bzl",
    "declare_out_files",
    "get_include_directory",
    "get_out_dir",
    "get_plugin_args",
    "get_proto_arguments",
    "includes_from_deps",
    "protos_from_context",
)

_GENERATED_PROTO_FORMAT = "{}_pb2.py"
_GENERATED_GRPC_PROTO_FORMAT = "{}_pb2_grpc.py"

def _generate_py_impl(context):
    protos = protos_from_context(context)
    includes = includes_from_deps(context.attr.deps)
    out_files = declare_out_files(protos, context, _GENERATED_PROTO_FORMAT)
    tools = [context.executable._protoc]

    out_dir = get_out_dir(protos, context)
    arguments = ([
        "--python_out={}".format(out_dir.path),
    ] + [
        "--proto_path={}".format(get_include_directory(i))
        for i in includes
    ] + [
        "--proto_path={}".format(context.genfiles_dir.path),
    ])
    if context.attr.plugin:
        arguments += get_plugin_args(
            context.executable.plugin,
            [],
            out_dir.path,
            False,
            context.attr.plugin.label.name,
        )
        tools.append(context.executable.plugin)

    arguments += get_proto_arguments(protos, context.genfiles_dir.path)

    context.actions.run(
        inputs = protos + includes,
        tools = tools,
        outputs = out_files,
        executable = context.executable._protoc,
        arguments = arguments,
        mnemonic = "ProtocInvocation",
    )

    imports = []
    if out_dir.import_path:
        imports.append("%s/%s/%s" % (context.workspace_name, context.label.package, out_dir.import_path))

    return [
        DefaultInfo(files = depset(direct = out_files)),
        PyInfo(
            transitive_sources = depset(),
            imports = depset(direct = imports),
        ),
    ]

_generate_pb2_src = rule(
    attrs = {
        "deps": attr.label_list(
            mandatory = True,
            allow_empty = False,
            providers = [ProtoInfo],
        ),
        "plugin": attr.label(
            mandatory = False,
            executable = True,
            providers = ["files_to_run"],
            cfg = "host",
        ),
        "_protoc": attr.label(
            default = Label("//external:protocol_compiler"),
            providers = ["files_to_run"],
            executable = True,
            cfg = "host",
        ),
    },
    implementation = _generate_py_impl,
)

def py_proto_library(
        name,
        deps,
        plugin = None,
        **kwargs):
    """Generate python code for a protobuf.

    Args:
      name: The name of the target.
      deps: A list of proto_library dependencies. Must contain a single element.
      plugin: An optional custom protoc plugin to execute together with
        generating the protobuf code.
      **kwargs: Additional arguments to be supplied to the invocation of
        py_library.
    """
    codegen_target = "_{}_codegen".format(name)
    if len(deps) != 1:
        fail("Can only compile a single proto at a time.")

    _generate_pb2_src(
        name = codegen_target,
        deps = deps,
        plugin = plugin,
        **kwargs
    )

    native.py_library(
        name = name,
        srcs = [":{}".format(codegen_target)],
        deps = [
            "@com_google_protobuf//:protobuf_python",
            ":{}".format(codegen_target),
        ],
        **kwargs
    )

def _generate_pb2_grpc_src_impl(context):
    protos = protos_from_context(context)
    includes = includes_from_deps(context.attr.deps)
    out_files = declare_out_files(protos, context, _GENERATED_GRPC_PROTO_FORMAT)

    plugin_flags = ["grpc_2_0"] + context.attr.strip_prefixes

    arguments = []
    tools = [context.executable._protoc, context.executable._grpc_plugin]
    out_dir = get_out_dir(protos, context)
    arguments += get_plugin_args(
        context.executable._grpc_plugin,
        plugin_flags,
        out_dir.path,
        False,
    )
    if context.attr.plugin:
        arguments += get_plugin_args(
            context.executable.plugin,
            [],
            out_dir.path,
            False,
            context.attr.plugin.label.name,
        )
        tools.append(context.executable.plugin)

    arguments += [
        "--proto_path={}".format(get_include_directory(i))
        for i in includes
    ]
    arguments += ["--proto_path={}".format(context.genfiles_dir.path)]
    arguments += get_proto_arguments(protos, context.genfiles_dir.path)

    context.actions.run(
        inputs = protos + includes,
        tools = tools,
        outputs = out_files,
        executable = context.executable._protoc,
        arguments = arguments,
        mnemonic = "ProtocInvocation",
    )

    return [
        DefaultInfo(files = depset(direct = out_files)),
        PyInfo(
            transitive_sources = depset(),
            # Imports are already configured by the generated py impl
            imports = depset(),
        ),
    ]

_generate_pb2_grpc_src = rule(
    attrs = {
        "deps": attr.label_list(
            mandatory = True,
            allow_empty = False,
            providers = [ProtoInfo],
        ),
        "strip_prefixes": attr.string_list(),
        "plugin": attr.label(
            mandatory = False,
            executable = True,
            providers = ["files_to_run"],
            cfg = "host",
        ),
        "_grpc_plugin": attr.label(
            executable = True,
            providers = ["files_to_run"],
            cfg = "host",
            default = Label("//src/compiler:grpc_python_plugin"),
        ),
        "_protoc": attr.label(
            executable = True,
            providers = ["files_to_run"],
            cfg = "host",
            default = Label("//external:protocol_compiler"),
        ),
    },
    implementation = _generate_pb2_grpc_src_impl,
)

def py_grpc_library(
        name,
        srcs,
        deps,
        plugin = None,
        strip_prefixes = [],
        **kwargs):
    """Generate python code for gRPC services defined in a protobuf.

    Args:
      name: The name of the target.
      srcs: (List of `labels`) a single proto_library target containing the
        schema of the service.
      deps: (List of `labels`) a single py_proto_library target for the
        proto_library in `srcs`.
      strip_prefixes: (List of `strings`) If provided, this prefix will be
        stripped from the beginning of foo_pb2 modules imported by the
        generated stubs. This is useful in combination with the `imports`
        attribute of the `py_library` rule.
      plugin: An optional custom protoc plugin to execute together with
        generating the gRPC code.
      **kwargs: Additional arguments to be supplied to the invocation of
        py_library.
    """
    codegen_grpc_target = "_{}_grpc_codegen".format(name)
    if len(srcs) != 1:
        fail("Can only compile a single proto at a time.")

    if len(deps) != 1:
        fail("Deps must have length 1.")

    _generate_pb2_grpc_src(
        name = codegen_grpc_target,
        deps = srcs,
        strip_prefixes = strip_prefixes,
        plugin = plugin,
        **kwargs
    )

    native.py_library(
        name = name,
        srcs = [
            ":{}".format(codegen_grpc_target),
        ],
        deps = [
            Label("//src/python/grpcio/grpc:grpcio"),
        ] + deps + [
            ":{}".format(codegen_grpc_target),
        ],
        **kwargs
    )

def py2and3_test(
        name,
        py_test = native.py_test,
        **kwargs):
    """Runs a Python test under both Python 2 and Python 3.

    Args:
      name: The name of the test.
      py_test: The rule to use for each test.
      **kwargs: Keyword arguments passed directly to the underlying py_test
        rule.
    """
    if "python_version" in kwargs:
        fail("Cannot specify 'python_version' in py2and3_test.")

    names = [name + suffix for suffix in (".python2", ".python3")]
    python_versions = ["PY2", "PY3"]
    for case_name, python_version in zip(names, python_versions):
        py_test(
            name = case_name,
            python_version = python_version,
            **kwargs
        )

    suite_kwargs = {}
    if "visibility" in kwargs:
        suite_kwargs["visibility"] = kwargs["visibility"]

    native.test_suite(
        name = name,
        tests = names,
        **suite_kwargs
    )