# Copyright 2025 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.
"""A rule defining importable Python modules that simplify locating runfiles."""

load("@rules_python//python:defs.bzl", "py_library")

_TEMPLATE = """# FILE GENERATED BY {label}, DO NOT EDIT!

from pw_build.python_runfiles import PythonRunfilesLabelAdapter

RLOCATION = PythonRunfilesLabelAdapter(
    runfiles_path = "{runfiles_path}",
    source_repo = {source_repo},
)
"""

def _generated_runfile_import_impl(ctx):
    import_path = "/".join(
        [ctx.attr.import_dir] + ctx.attr.import_location.split("."),
    )

    generated_import = ctx.actions.declare_file(import_path + ".py")

    if ctx.file.src:
        f = ctx.file.src
        runfiles = ctx.runfiles(ctx.files.src)
        runfiles = runfiles.merge(ctx.attr.src[DefaultInfo].data_runfiles)
        target_repo_name = ctx.attr.src.label.repo_name
    elif ctx.executable.bin:
        f = ctx.executable.bin
        runfiles = ctx.runfiles()
        runfiles = runfiles.merge(ctx.attr.bin[DefaultInfo].data_runfiles)
        target_repo_name = ctx.attr.bin.label.repo_name
    else:
        fail(ctx.label, "requires `src` to be set")

    if not target_repo_name:
        target_repo_name = ctx.attr.module_name

    # It's valid for this to be None. For example, the repo_name of the
    # root repo defaults to None to support the canonical name of "@@".
    current_repo_name = ctx.label.repo_name
    current_repo_name = '"{}"'.format(current_repo_name) if current_repo_name else None

    runfile_path = f.short_path

    # External runfiles are of the form `../+_repo_blah/`, which needs to be
    # stripped. The `..` is unneeded, and the `+_repo_blah` is redundant.
    # This is ugly, but rules_python has to do something similar.
    if runfile_path.startswith(".."):
        runfile_path = "/".join(runfile_path.split("/")[2:])
    runfile_path = "{}/{}".format(target_repo_name, runfile_path)

    ctx.actions.write(
        output = generated_import,
        content = _TEMPLATE.format(
            label = ctx.label,
            runfiles_path = runfile_path,
            source_repo = current_repo_name,
        ),
    )

    return DefaultInfo(
        files = depset(direct = [generated_import]),
        runfiles = runfiles,
    )

_generated_runfile_import = rule(
    implementation = _generated_runfile_import_impl,
    attrs = {
        "bin": attr.label(executable = True, cfg = "target", allow_files = True),
        "import_dir": attr.string(mandatory = True),
        "import_location": attr.string(mandatory = True),
        "module_name": attr.string(),
        "src": attr.label(allow_single_file = True),
    },
)

def pw_py_importable_runfile(*, name, src = None, executable = False, import_location = None, **kwargs):
    """An importable py_library that makes loading runfiles easier.

    When using Bazel runfiles from Python, ``Rlocation()`` takes two arguments:

    1. The ``path`` of the runfiles. This is the apparent repo name joined with
       the path within that repo.
    2. The ``source_repo`` to evaluate ``path`` from. This is related to how
       apparent repo names and canonical repo names are handled by Bazel.

    Unfortunately, it's easy to get these arguments wrong.

    This generated Python library short-circuits this problem by letting Bazel
    generate the correct arguments to ``Rlocation()`` so users don't even have
    to think about what to pass.

    For example:

        ```
        # In @bloaty//:BUILD.bazel, or wherever is convenient:
        pw_py_importable_runfile(
            name = "bloaty_runfiles",
            src = "//:bin/bloaty",
            executable = True,
            import_location = "bloaty.bloaty_binary",
        )

        # Using the pw_py_importable_runfile from a py_binary in a
        # BUILD.bazel file:
        py_binary(
            name = "my_binary",
            srcs = ["my_binary.py"],
            main = "my_binary.py",
            deps = ["@bloaty//:bloaty_runfiles"],
        )

        # In my_binary.py:
        import bloaty.bloaty_binary
        from python.runfiles import runfiles  # type: ignore

        r = runfiles.Create()
        bloaty_path = r.Rlocation(*bloaty.bloaty_binary.RLOCATION)
        ```

    Note: Because this exposes runfiles as importable Python modules,
    the import paths of the generated libraries may collide with existing
    Python libraries. When this occurs, you need to extend the import path
    of modules with generated files.

    Args:
        name: name of the target.
        import_location: The final Python import path of the generated module.
            By default, this is ``path.to.package.label_name``.
        src: The file this library exposes as runfiles.
        executable: Whether or not the source file is executable.
        **kwargs: Common attributes to forward both underlying targets.
    """
    _generated_py_file = "{}._generated_py_import".format(name)
    _virtual_import_dir = "_{}_virtual_imports".format((name))
    if not import_location:
        import_location = "/".join((
            native.package_relative_label(":{}".format(name)).package,
            name,
        ))
    _generated_runfile_import(
        name = _generated_py_file,
        src = src if not executable else None,
        bin = src if executable else None,
        import_location = import_location,
        import_dir = _virtual_import_dir,
        module_name = native.module_name(),
        **kwargs
    )
    py_library(
        name = name,
        imports = [_virtual_import_dir],
        srcs = [":" + _generated_py_file],
        data = [src],
        deps = [
            Label("//pw_build/py:python_runfiles"),
            Label("@rules_python//python/runfiles"),
        ],
        **kwargs
    )