load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
load("@capnp-cpp//src/capnp:cc_capnp_library.bzl", "cc_capnp_library")

CAPNP_TEMPLATE = """@{schema_id};

# generated by @workerd//build/wd_js_bundle.bzl

using Modules = import "/workerd/jsg/modules.capnp";

const {const_name} :Modules.Bundle = (
  modules = [
{modules}
]);
"""

def _to_name(file_name):
    return file_name.removesuffix(".js")

def _to_d_ts(file_name):
    return file_name.removesuffix(".js") + ".d.ts"

def _relative_path(file_path, dir_path):
    if not file_path.startswith(dir_path):
        fail("file_path need to start with dir_path: " + file_path + " vs " + dir_path)
    return file_path.removeprefix(dir_path)

def _gen_compile_cache_impl(ctx):
    file_list = ctx.actions.declare_file("in")

    srcs = []
    for src in ctx.attr.srcs:
        srcs.extend(src.files.to_list())

    outs = [ctx.actions.declare_file(src.basename + "_cache") for src in srcs]

    content = []
    for i in range(0, len(srcs)):
        content.append("{} {}".format(srcs[i].path, outs[i].path))

    ctx.actions.write(
        output = file_list,
        content = "\n".join(content) + "\n",
    )

    args = ctx.actions.args()
    args.add(file_list)

    run_under = ctx.attr._run_under[BuildSettingInfo].value

    # use run_shell together with cfg = target instead of ctx.actions.run
    # to prevent double-compilation of v8.
    ctx.actions.run_shell(
        outputs = outs,
        inputs = [file_list] + srcs,
        command = run_under + " " + ctx.executable._tool.path + " $@",
        arguments = [args],
        use_default_shell_env = True,
        tools = [ctx.executable._tool],
    )

    return [
        DefaultInfo(files = depset(direct = outs)),
    ]

_gen_compile_cache = rule(
    implementation = _gen_compile_cache_impl,
    attrs = {
        "srcs": attr.label_list(mandatory = True, allow_files = True),
        "_tool": attr.label(
            executable = True,
            allow_single_file = True,
            cfg = "target",
            default = "//src/rust/gen-compile-cache",
        ),
        "_run_under": attr.label(default = "//build/config:target_run_under"),
    },
)

def _get_compile_cache(compile_cache, m):
    if not compile_cache:
        return None
    files = m.files.to_list()

    if len(files) != 1:
        fail("only single file expected")

    return compile_cache.get(files[0].path)

MODULE_TEMPLATE = """    (name = "{name}", {src_type} = embed "{path}", type = {type}, {extras})"""

def _gen_api_bundle_capnpn_impl(ctx):
    output_dir = ctx.outputs.out.dirname + "/"

    def _render_module(name, label, src_type, type, cache = None):
        ts_declaration_extra = (
            "tsDeclaration = embed \"" + _relative_path(
                ctx.expand_location("$(location {})".format(ctx.attr.declarations[name]), ctx.attr.data),
                output_dir,
            ) + "\", "
        ) if name in ctx.attr.declarations else ""
        cache_extra = (
            "compileCache = embed \"{}\", ".format(_relative_path(cache, output_dir))
        ) if cache else ""
        extras = ts_declaration_extra + cache_extra

        return MODULE_TEMPLATE.format(
            name = name,
            # capnp doesn't allow ".." dir escape, make paths relative.
            # this won't work for embedding paths outside of rule directory subtree.
            src_type = src_type,
            path = _relative_path(
                ctx.expand_location("$(location {})".format(label), ctx.attr.data),
                output_dir,
            ),
            type = type,
            extras = extras,
        )

    compile_cache = {}
    if ctx.attr.compile_cache:
        locations = ctx.expand_location("$(locations {})".format(ctx.attr.compile_cache.label)).split(" ")
        for loc in locations:
            compile_cache[loc.removesuffix("_cache")] = loc

    modules = [
        _render_module(ctx.attr.builtin_modules[m], m.label, "src", "builtin", _get_compile_cache(compile_cache, m))
        for m in ctx.attr.builtin_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_modules[m], m.label, "src", "internal", _get_compile_cache(compile_cache, m))
        for m in ctx.attr.internal_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_wasm_modules[m], m.label, "wasm", "internal")
        for m in ctx.attr.internal_wasm_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_data_modules[m], m.label, "data", "internal")
        for m in ctx.attr.internal_data_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_json_modules[m], m.label, "json", "internal")
        for m in ctx.attr.internal_json_modules
    ]

    content = CAPNP_TEMPLATE.format(
        schema_id = ctx.attr.schema_id,
        modules = ",\n".join(modules),
        const_name = ctx.attr.const_name,
    )
    ctx.actions.write(ctx.outputs.out, content)

gen_api_bundle_capnpn = rule(
    implementation = _gen_api_bundle_capnpn_impl,
    attrs = {
        "schema_id": attr.string(mandatory = True),
        "out": attr.output(mandatory = True),
        "builtin_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_wasm_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_data_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_json_modules": attr.label_keyed_string_dict(allow_files = True),
        "declarations": attr.string_dict(),
        "data": attr.label_list(allow_files = True),
        "const_name": attr.string(mandatory = True),
        "deps": attr.label_list(),
        "compile_cache": attr.label(),
    },
)

def _copy_modules(modules, declarations, out_dir = ""):
    """Copy files from the modules map to the current package.

    Returns new module map using file copies.
    This is necessary since capnp compiler doesn't allow embeds outside of current subidrectory.
    """
    result = dict()
    declarations_result = dict()
    for m in modules:
        new_filename = out_dir + modules[m].replace(":", "_").replace("/", "_")
        copy_file(name = new_filename + "@copy", src = m, out = new_filename)

        m_d_ts = _to_d_ts(m)
        if m_d_ts in declarations:
            new_d_ts_filename = new_filename + ".d.ts"
            copy_file(name = new_d_ts_filename + "@copy", src = m_d_ts, out = new_d_ts_filename)
            declarations_result[modules[m]] = str(native.package_relative_label(new_d_ts_filename))

        result[new_filename] = modules[m]
    return result, declarations_result

def wd_js_bundle(
        name,
        import_name,
        schema_id,
        builtin_modules,
        internal_modules = [],
        internal_wasm_modules = [],
        internal_data_modules = [],
        internal_json_modules = [],
        declarations = [],
        deps = [],
        gen_compile_cache = False,
        out_dir = ""):
    """Generate cc capnp library with js api bundle.

    NOTE: Due to capnpc embed limitation all modules must be in the same or sub directory of the
          actual rule usage.

    Args:
     name: cc_capnp_library rule name
     import_name: The js import specifier. builtin modules are accessible from
                  user code under `<import_name>:<module_name>`, internal modules
                  are accessible from builtin modules under
                  `<import_name>-internal:<module_name>`
                  The capnproto bundle object generated will be called
                  `import_name` + "Bundle"
     schema_id: capnpn schema id
     builtin_modules: list of js source files for builtin modules
     internal_modules: list of js source files for internal modules
     internal_wasm_modules: list of wasm source files
     internal_data_modules: list of data source files
     internal_json_modules: list of json source files
     declarations: d.ts label set
     deps: dependency list
     gen_compile_cache: generate compilation cache of every file and include into the bundle
    """
    builtin_modules_dict = {
        m: "{}:{}".format(import_name, _to_name(m).removeprefix(out_dir))
        for m in builtin_modules
    }
    internal_modules_dict = {
        m: "{}-internal:{}".format(import_name, _to_name(m.removeprefix(out_dir).removeprefix("internal/")))
        for m in internal_modules
    }
    internal_wasm_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix(out_dir).removeprefix("internal/"))
        for m in internal_wasm_modules
    }
    internal_data_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix(out_dir).removeprefix("internal/"))
        for m in internal_data_modules
    }
    internal_json_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix(out_dir).removeprefix("internal/"))
        for m in internal_json_modules
    }

    builtin_modules_dict, builtin_declarations = _copy_modules(
        builtin_modules_dict,
        declarations,
        out_dir,
    )
    internal_modules_dict, internal_declarations = _copy_modules(
        internal_modules_dict,
        declarations,
        out_dir,
    )
    internal_wasm_modules_dict, _ = _copy_modules(
        internal_wasm_modules_dict,
        declarations,
        out_dir,
    )
    internal_data_modules_dict, _ = _copy_modules(
        internal_data_modules_dict,
        declarations,
        out_dir,
    )
    internal_json_modules_dict, _ = _copy_modules(
        internal_json_modules_dict,
        declarations,
        out_dir,
    )

    data = (
        list(builtin_modules_dict) +
        list(internal_modules_dict) +
        list(internal_wasm_modules_dict) +
        list(internal_data_modules_dict) +
        list(internal_json_modules_dict) +
        list(builtin_declarations.values()) +
        list(internal_declarations.values())
    )

    compile_cache = None
    if gen_compile_cache:
        _gen_compile_cache(
            name = name + "@compile_cache",
            srcs = builtin_modules_dict.keys() + internal_modules_dict.keys(),
        )
        compile_cache = name + "@compile_cache"
        deps = deps + [compile_cache]
        data = data + [compile_cache]

    gen_api_bundle_capnpn(
        name = name + "@gen",
        out = name + ".capnp",
        schema_id = schema_id,
        const_name = import_name + "Bundle",
        builtin_modules = builtin_modules_dict,
        internal_modules = internal_modules_dict,
        internal_wasm_modules = internal_wasm_modules_dict,
        internal_data_modules = internal_data_modules_dict,
        internal_json_modules = internal_json_modules_dict,
        declarations = builtin_declarations | internal_declarations,
        data = data,
        deps = deps,
        compile_cache = compile_cache,
    )

    cc_capnp_library(
        name = name,
        srcs = [name + ".capnp"],
        strip_include_prefix = "",
        visibility = ["//visibility:public"],
        data = data,
        deps = ["@workerd//src/workerd/jsg:modules_capnp"],
        include_prefix = import_name,
    )