# Copyright 2023 The Bazel Authors. All rights reserved.
#
# 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
#
#    http://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.

"""
Defines baseline profiles processing methods in Android Rules.
"""

_BASELINE_PROFILE_DIR_SUFFIX = "-baseline-profile/"

def _process(
        ctx,
        transitive_profiles = depset(),
        startup_profiles = [],
        deploy_jar = None,
        has_proguard_specs = False,
        enable_optimizer_integration = False,
        merge_tool = None,
        profgen = None,
        toolchain_type = None):
    """Processes all the transitive baseline profiles.

    Baseline profiles propagated from libraries will be merged, and if optimizer integration is
    enabled, startup profiles will be merged as well, and wildcards in baseline profiles will be
    expanded.

    Args:
      ctx: The context.
      transitive_profiles: Depset. The transitive baseline profiles propagated from android_library
          and aar_import.
      startup_profiles: List. The startup profiles.
      deploy_jar: File. The deploy jar.
      has_proguard_specs: Boolean. Whether to have proguard specs.
      enable_optimizer_integration: Boolean. Whether to use startup profile and baseline profiles in optimization.
      merge_tool: FilesToRunProvider. An executable for merging baseline profiles.
      profgen: FilesToRunProvider: An executable for compiling baseline profiles.
      toolchain_type: Label or String. Toolchain type of the executable used in actions.
    Returns:
      A struct containing all the outputs from processing baseline profiles.
    """

    baseline_profile = None
    startup_profile = None
    if transitive_profiles:
        baseline_profile = _get_profile_artifact(ctx, "static-prof.txt")
        _merge(
            ctx,
            baseline_profile,
            transitive_profiles,
            "MergeBaselineProfiles",
            merge_tool = merge_tool,
            toolchain_type = toolchain_type,
        )

    if has_proguard_specs and enable_optimizer_integration:
        # This is only needed for optimized builds since otherwise the dexer doesn't process this.
        if startup_profiles:
            startup_profile = _get_profile_artifact(ctx, "static-startup-prof.txt")
            _merge(
                ctx,
                output = startup_profile,
                inputs = ctx.files.startup_profiles,
                mnemonic = "MergeStartupProfiles",
                merge_tool = merge_tool,
                toolchain_type = toolchain_type,
            )

        # Wildcards only need to be expanded for optimized builds since if these aren't consumed by
        # the optimizer, they can just be expanded during profile compilation instead.
        # Start-up profiles are not expanded because it shouldn't be necessary as these should
        # contain profiles generated by devices on start-up.
        if baseline_profile:
            expanded_baseline_profile = _get_profile_artifact(ctx, "expanded-static-prof.txt")
            _expand_wildcards(
                ctx,
                output = expanded_baseline_profile,
                deploy_jar = deploy_jar,
                profile = baseline_profile,
                profgen = profgen,
                toolchain_type = toolchain_type,
            )
            baseline_profile = expanded_baseline_profile
    return struct(
        baseline_profile = baseline_profile,
        startup_profile = startup_profile,
    )

def _process_art_profile(
        ctx,
        final_classes_dex,
        merged_profile,
        proguard_output_map = None,
        profgen = None,
        zipper = None,
        toolchain_type = None):
    """ Compiles the merged baseline profile.

    Profiles are compiled with profgen into binary ART profiles. The binary
    profiles will be bundled into the final APK and used at installation time to speed up app
    startup and reduce jank.

    Args:
      ctx: The context.
      final_classes_dex: File. Final classes zip artifact.
      merged_profile: File. The merged profile from transitive baseline profile files.
      proguard_output_map: File. Optional. The proguard output mapping file.
      profgen: FilesToRunProvider. The profgen executable for profile compilation.
      zipper: FilesToRunProvider. An executable for compressing files to an archive.
      toolchain_type: Label or String. Toolchain type of the executable used in actions.
    Returns:
      Provider info containing BaselineProfileProvider for all merged profiles.
    """

    # Profgen
    output_profile = _get_profile_artifact(ctx, "baseline.prof")
    output_profile_meta = _get_profile_artifact(ctx, "baseline.profm")
    profgen_inputs = [final_classes_dex, merged_profile]
    profgen_args = ctx.actions.args()
    profgen_args.add("bin", merged_profile)
    profgen_args.add("--apk", final_classes_dex)
    profgen_args.add("--output", output_profile)
    profgen_args.add("--output-meta", output_profile_meta)
    if proguard_output_map:
        profgen_args.add("--map", proguard_output_map)
        profgen_inputs.append(proguard_output_map)
    ctx.actions.run(
        mnemonic = "GenerateARTProfile",
        executable = profgen,
        progress_message = "Generating Android P-R ART profile for %{label} APK",
        arguments = [profgen_args],
        inputs = profgen_inputs,
        outputs = [output_profile, output_profile_meta],
        use_default_shell_env = True,
        toolchain = toolchain_type,
    )

    # Zip ART profiles
    output_profile_zip = _get_profile_artifact(ctx, "art_profile.zip")
    zip_args = ctx.actions.args()
    zip_args.add("c", output_profile_zip)
    zip_args.add(output_profile.path, format = "assets/dexopt/baseline.prof=%s")
    zip_args.add(output_profile_meta.path, format = "assets/dexopt/baseline.profm=%s")
    ctx.actions.run(
        mnemonic = "ZipARTProfiles",
        executable = zipper,
        progress_message = "Zip ART Profiles for %{label}",
        arguments = [zip_args],
        inputs = [output_profile, output_profile_meta],
        outputs = [output_profile_zip],
        use_default_shell_env = True,
        toolchain = toolchain_type,
    )
    return BaselineProfileProvider(
        # Unnecessary to pass the transitive profiles to native rule
        depset(),
        output_profile_zip,
    )

def _get_profile_dir(ctx):
    return ctx.label.name + _BASELINE_PROFILE_DIR_SUFFIX

def _get_profile_artifact(ctx, name):
    return ctx.actions.declare_file(_get_profile_dir(ctx) + name)

def _merge(
        ctx,
        output,
        inputs = [],
        mnemonic = "MergeBaselineProfiles",
        merge_tool = None,
        toolchain_type = None):
    args = ctx.actions.args()
    args.add_all(inputs, before_each = "--input")
    args.add("--output", output)

    ctx.actions.run(
        executable = merge_tool,
        mnemonic = mnemonic,
        arguments = [args],
        inputs = inputs,
        outputs = [output],
        use_default_shell_env = True,
        toolchain = toolchain_type,
    )

def _expand_wildcards(
        ctx,
        output,
        deploy_jar = None,
        profile = None,
        profgen = None,
        toolchain_type = None):
    args = ctx.actions.args()
    args.add("expandWildcards", deploy_jar)
    args.add("--profile", profile)
    args.add("--output", output)

    ctx.actions.run(
        executable = profgen,
        outputs = [output],
        inputs = [deploy_jar, profile],
        arguments = [args],
        mnemonic = "ExpandBaselineProfileWildcards",
        progress_message = "Expanding baseline profile wildcards for %{label} APK",
        toolchain = toolchain_type,
    )

baseline_profiles = struct(
    expand_wildcards = _expand_wildcards,
    get_profile_artifact = _get_profile_artifact,
    process = _process,
    process_art_profile = _process_art_profile,
)