#!/usr/bin/env lucicfg
# Copyright 2020 The Flutter Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""
Utility for loading Flutter infrastructure's yaml config into LUCI.

See https://github.com/flutter/cocoon/blob/main/CI_YAML.md for more information.
"""

proto.new_descriptor_set(
    name = "scheduler",
    blob = io.read_file("./scheduler.bin"),
).register()

load("@proto//scheduler.proto", _scheduler_pb = "scheduler")
load("//lib/common.star", "common")
load("//lib/copy.star", "copy")
load("//lib/helpers.star", "helpers")
load("//lib/repos.star", "repos")

def _ci_yaml(repo, branch):
    """Creates ci_yaml object.

    ci_yamls are expected to have a generated jspb file under
    //config/generated/ci_yaml with the format `$repo_config.json` for tip of
    tree or `$repo_$branch_config.json` for release branches.

    ci_yaml json files are autogenerated based on the .ci.yaml in $repo.

    Args:
        repo(str): Name of the repository to follow.
        branch(str): Branch name.
    """
    version = "" if _is_default_branch(branch) else "_%s" % branch
    return proto.from_jsonpb(
        _scheduler_pb.SchedulerConfig,
        io.read_file("./generated/ci_yaml/%s%s_config.json" % (repo, version)),
    )

# lucicfg does not currently support enums from protos.
# To work around this, statically redeclare the values from:
#   //config/lib/ci_yaml/scheduler.proto - SchedulerSytem
SCHEDULER_COCOON = 1
SCHEDULER_LUCI = 2

# lucicfg is only intended to generate LUCI based targets.
SUPPORTED_SCHEDULERS = [SCHEDULER_COCOON, SCHEDULER_LUCI]

def _recipes(repo, branch, recipes_ref):
    """Return the set of recipes specified in ci.yaml."""
    recipe_names = {}
    ci_yaml = _ci_yaml(repo, branch)
    for target in ci_yaml.targets:
        # Only LUCI based targets have recipes
        if target.scheduler not in SUPPORTED_SCHEDULERS:
            continue
        if target.enabled_branches and branch not in target.enabled_branches:
            continue
        recipe_name = _full_recipe_name(target.recipe, recipes_ref)
        if recipe_name not in recipe_names.keys():
            recipe_names[recipe_name] = target.recipe
    for recipe_name, recipe in recipe_names.items():
        luci.recipe(
            name = recipe_name,
            recipe = recipe,
            cipd_package = "flutter/recipe_bundles/flutter.googlesource.com/recipes",
            cipd_version = recipes_ref,
            use_bbagent = True,
        )

def _add_recipes_cq(target):
    if "add_recipes_cq" in target.properties:
        return target.properties["add_recipes_cq"] == "true"
    return False

def _swarming_caches(repo, version, dependencies = [], ios_runtimes = [], xcode = None):
    """Return the equivalent swarming caches for target.

    Caches are generated based on a targets dependencies, the version, and their branch.
    """
    if version == None:
        version = "main"
    legacy_paths = {
        "android_sdk": "android",
        "chrome_and_driver": "chrome",
        "open_jdk": "java",
    }
    swarming_caches = [
        swarming.cache(path = "builder", name = repo + "_" + version + "_builder"),
        swarming.cache(path = ".pub-cache", name = "pub_cache"),
    ]
    for dependency in dependencies:
        path = dependency["dependency"]
        if path == "xcode":
            continue
        dep_version = dependency["version"] if "version" in dependency else "default"
        name = repo + "_" + version + "_" + path + "_" + dep_version
        name = name.replace(":", "_")
        name = name.replace("-", "_")
        name = name.replace(".", "_")
        swarming_caches.append(swarming.cache(name = name, path = path))
        if path in legacy_paths:
            swarming_caches.append(swarming.cache(name = name + "_legacy", path = legacy_paths[path]))
    for runtime in ios_runtimes:
        path = "xcode_runtime_" + runtime
        path = path.replace("-", "_")
        name = repo + "_" + version + "_" + path
        swarming_caches.append(swarming.cache(name = name, path = path))
    if xcode:
        name = repo + "_" + version + "_xcode_" + xcode
        swarming_caches.append(swarming.cache(name = name + "_legacy", path = "osx_sdk"))
        swarming_caches.append(swarming.cache(name = name, path = "xcode"))
    return swarming_caches

def _legacy_swarming_caches(ci_yaml, target, platform):
    """DEPRECATED: Return the swarming caches based on what is listed in ci.yaml."""
    swarming_caches = []
    caches = []
    cache_set = []

    # Target specific caches
    if target and "caches" in target.properties:
        caches = json.decode(target.properties["caches"])

        # Mark all in unique set to ensure platform does not override
        for cache in caches:
            cache_set.append(cache["name"])

    # Platform wide caches
    if "caches" in ci_yaml.platform_properties[platform].properties:
        platform_caches = json.decode(ci_yaml.platform_properties[platform].properties["caches"])

        # Ensure platform wide caches do not override target level caches
        for platform_cache in platform_caches:
            if platform_cache["name"] not in cache_set:
                caches.append(platform_cache)

    # Generate LUCI config caches
    for cache in caches:
        swarming_caches.append(swarming.cache(name = cache["name"], path = cache["path"]))
    return swarming_caches

def _platform_dimensions(ci_yaml):
    """Gets platform_dimensions from ci_yaml."""
    platform_dimensions = {}
    for platform, _ in dict(ci_yaml.platform_properties).items():
        platform_dimensions[platform] = {}
        for k, v in dict(ci_yaml.platform_properties[platform].dimensions).items():
            platform_dimensions[platform][k] = _parse_attribute(k, v)
    return platform_dimensions

def _platform_properties(ci_yaml):
    """Gets platform_properties from ci_yaml."""
    platform_properties = {}
    for platform, _ in dict(ci_yaml.platform_properties).items():
        platform_properties[platform] = {}
        for k, v in dict(ci_yaml.platform_properties[platform].properties).items():
            platform_properties[platform][k] = _parse_attribute(k, v)
    return platform_properties

def _parse_attribute(key, value):
    """Helper function to parse non string type properties."""
    json_list = ["android_sdk_license", "android_sdk_preview_license"]
    if value == "true":
        return True
    elif value == "false":
        return False
    elif key in json_list:
        # Yaml -> JSON -> Here adds an extra string escape that we unescape here
        return value.replace("\\n", "\n")
    elif value.startswith("["):
        return json.decode(value)
    elif value.isdigit():
        return int(value)
    else:
        return value

def _dimensions(ci_yaml, target):
    dimensions = {}

    # Update with platform specific dimensions
    platform = _target_platform(target)
    input_platform_dimensions = _platform_dimensions(ci_yaml)
    if platform in input_platform_dimensions:
        for k, v in input_platform_dimensions[platform].items():
            dimensions[k] = v

    # Update with target specific dimensions
    for k, v in dict(target.dimensions).items():
        dimensions[k] = _parse_attribute(k, v)
    return dimensions

def _properties(ci_yaml, target, default_properties = {}, repo = None, branch = None):
    """Creates builder properties based on ci_yaml.

    Args:
        ci_yaml(SchedulerConfig): Config file that contains target.
        target(Target): Configuration to generate properties for.
        default_properties(dict): Repo level properties.
        repo(string): Name of the repo.
        branch(string): Name of the branch.
    """
    properties = {}

    # Initialize with default repo properties
    for k, v in default_properties.items():
        properties[k] = v

    # Update with platform specific properties
    platform = _target_platform(target)
    input_platform_properties = _platform_properties(ci_yaml)
    if platform in input_platform_properties:
        for k, v in input_platform_properties[platform].items():
            properties[k] = v

    # Finally, update with target specific properties
    for k, v in dict(target.properties).items():
        if k == "dependencies":
            # Merge target dependencies with platform
            dependencies = json.decode(v)
            target_dependency_lookup_table = []
            for dep in dependencies:
                target_dependency_lookup_table.append(dep["dependency"])

            # Add platform dependencies that are not already defined
            for dep in properties[k]:
                if dep["dependency"] not in target_dependency_lookup_table:
                    dependencies.append(dep)
            properties[k] = dependencies
        else:
            properties[k] = _parse_attribute(k, v)

    # Special case to pass xcode for engine
    for xcode_platform in ["mac", "mac_ios", "mac_arm64_ios"]:
        if platform != xcode_platform:
            continue
        if "xcode" not in properties:
            continue
        xcode_version = {
            "sdk_version": properties["xcode"],
        }

        # DeviceLab bots are on a slower lab network. We avoid LUCI caches on those bots
        if platform in ["mac_ios", "mac_arm64_ios"]:
            properties["$flutter/devicelab_osx_sdk"] = xcode_version
        else:
            properties["$flutter/osx_sdk"] = xcode_version

    # Special case to add runtimes for ios simultators
    if platform == "mac" and "runtime_versions" in properties:
        # Add to existing osx_sdk properties dict, or create new
        osx_sdk = properties.get("$flutter/osx_sdk", {})
        osx_sdk["runtime_versions"] = properties["runtime_versions"]
        properties["$flutter/osx_sdk"] = osx_sdk

    filter_list = ["caches"]
    for filter in filter_list:
        if filter in properties:
            properties.pop(filter)
    if repo:
        properties["git_repo"] = repo
    if branch:
        properties["git_branch"] = branch

    # Add `bringup` to properties.
    properties["bringup"] = target.bringup

    if "dependencies" not in properties:
        properties["dependencies"] = []
    return properties

def _full_recipe_name(recipe_name, recipes_ref):
    """Creates a recipe name for recipe and version.

    Args:
      recipe_name(str): This is a string with the recipe base name.
      recipes_ref(str): The git ref pointing to the recipes bundle to use.

    Returns:
      A string with the recipe's full name.
    """
    ref = recipes_ref.replace("refs/heads/", "")
    return "%s-%s" % (ref, recipe_name)

def _builder_name(builder_name, branch):
    """Creates a builder name that is versioned to branch.

    Args:
        builder_name(str): The original builder name.
        branch(str): Branch this builder is for.
    """
    if not _is_default_branch(branch):
        parts = builder_name.split(" ")
        parts.insert(1, branch)
        builder_name = " ".join(parts)
    return "%s|%s" % (builder_name, common.short_name(builder_name))

def _is_default_branch(branch):
    """Returns whether the branch is a default branch."""
    return branch in ("main", "master")

def _builders(repo, branch, version, testing_ref, dimensions, properties, triggering_policy, notifies = None, recipes_ref = "refs/heads/main"):
    """Return the builders specified in ci.yaml.

    Targets defined in .ci.yaml are translated to their LUCI equivalents.
    Mostly, this is forward fields to the LUCI equivalent, but there
    are some fields that are special cased.

    Args:
        repo(str): Name of the repo to create builders for.
        branch(str): Name of the branch to create builders for. Either None:dev|beta|stable.
        version(str): The branch number, following major_minor.
        testing_ref(str): The git ref to trigger builds on.
        dimensions(dict): Platform dimensions to pass.
        properties(dict): Repo level properties to pass to every builder.
        triggering_policy(scheduler): Scheduler policy to implement on postsubmit builds.
        notifies(list): List of luci.notifier to send notifications on events from builders.
        recipes_ref(str): The git ref pointing to the recipes bundle to use.
    """
    ci_yaml = _ci_yaml(repo, branch)

    # Only generate presubmit and staging builds on tip of tree
    staging_view_name = "%s_staging" % repo
    if _is_default_branch(branch):
        list_view_name = "%s-try" % repo
        luci.list_view(
            name = list_view_name,
            title = "%s try builders" % repo.capitalize(),
        )
        luci.console_view(
            name = staging_view_name,
            repo = repos.GIT_REMOTE[repo],
            refs = [testing_ref],
        )
    console_view_name = repo if _is_default_branch(branch) else "%s_%s" % (branch, repo)
    luci.console_view(
        name = console_view_name,
        repo = repos.GIT_REMOTE[repo],
        refs = [testing_ref],
    )

    # Defines prod schedulers
    trigger_name = "%s-gitiles-trigger-%s" % (branch, repo)
    luci.gitiles_poller(
        name = trigger_name,
        bucket = "prod",
        repo = repos.GIT_REMOTE[repo],
        refs = [testing_ref],
    )

    # Defines default triggering policy
    if not triggering_policy:
        triggering_policy = scheduler.greedy_batching(
            max_batch_size = 1,
            max_concurrent_invocations = 3,
        )

    for target in ci_yaml.targets:
        # Not all targets in ci.yaml are LUCI based, skip those.
        if target.scheduler not in SUPPORTED_SCHEDULERS:
            continue

        # Flaky tests should not be on releases
        if not _is_default_branch(branch) and target.bringup:
            continue
        platform = _target_platform(target)
        merged_dimensions = _dimensions(ci_yaml, target)
        merged_properties = _properties(ci_yaml, target, properties, repo, branch)
        ios_simulators = []
        if "runtime_versions" in merged_properties:
            ios_simulators = merged_properties["runtime_versions"]
        xcode = merged_properties["xcode"] if "xcode" in merged_properties else None
        caches = _swarming_caches(repo, version, merged_properties["dependencies"], ios_simulators, xcode)
        target_args = {
            "caches": caches,
            "dimensions": merged_dimensions,
            # On release branches, extend the timeout to handle goma misses.
            # TODO(godofredoc): Remove after hotfix 2.5.3 is released.
            "execution_timeout": target.timeout * time.minute * 3,
            "name": _builder_name(target.name, branch),
            "os": merged_properties["os"],
            "recipe": _full_recipe_name(target.recipe, recipes_ref),
            "repo": repos.GIT_REMOTE[repo],
            "properties": merged_properties,
        }

        # Set dimensions dimensions from properties.
        if platform in dimensions:
            target_args["dimensions"] = dimensions[platform]
        dimension_key_list = ("device_type", "device_os", "mac_model", "cores", "cpu")
        for dimension_key in dimension_key_list:
            if dimension_key in merged_properties:
                target_args["dimensions"][dimension_key] = str(merged_properties[dimension_key])

        # Try builds are generated for every target to enable led and presubmit runs
        if _is_default_branch(branch):
            presubmit_target_args = copy.deepcopy(target_args)
            presubmit_target_args["list_view_name"] = list_view_name
            if repo == "engine":
                presubmit_target_args["properties"]["no_lto"] = True
            if _add_recipes_cq(target):
                presubmit_target_args["add_cq"] = True
            common.common_try_builder(**presubmit_target_args)

        postsubmit_target_args = helpers.merge_dicts({
            "priority": 30,
            "notifies": notifies,
            "category": platform.capitalize(),
        }, target_args)

        # Postsubmit only generates if the target is enabled for this branch, which defaults to true.
        if not target.enabled_branches or branch in target.enabled_branches:
            # Set customized triggering policy for engine staging builders.
            # https://github.com/flutter/flutter/issues/92947
            if target.bringup:
                final_triggering_policy = scheduler.greedy_batching(
                    max_batch_size = 6,
                    max_concurrent_invocations = 1,
                )
            else:
                final_triggering_policy = triggering_policy

            # Mark triggering policy as none if the target is not scheduled by LUCI
            final_triggering_policy = final_triggering_policy if target.scheduler == SCHEDULER_LUCI else None

            # TODO(chillers): Remove when postsubmit properties are supported. https://github.com/flutter/flutter/issues/87675
            if target.recipe == "devicelab/devicelab_drone":
                if _is_default_branch(branch):
                    postsubmit_target_args["properties"]["upload_metrics"] = True

                    # Set a 2h expiration_timeout for non flaky devicelab builders.
                    if not target.bringup:
                        postsubmit_target_args["expiration_timeout"] = 120 * time.minute
            staging_target_args = helpers.merge_dicts({
                "bucket": "staging",
                "console_view_name": staging_view_name,
                "expiration_timeout": 24 * time.hour,
                "pool": "luci.flutter.staging",
            }, postsubmit_target_args)
            prod_target_args = helpers.merge_dicts({
                "bucket": "prod",
                "console_view_name": console_view_name,
                "pool": "luci.flutter.prod",
            }, postsubmit_target_args)

            # Only trigger targets in a single bucket based on bringup
            if target.bringup:
                staging_target_args["triggered_by"] = [trigger_name] if target.scheduler == SCHEDULER_LUCI else None
                staging_target_args["triggering_policy"] = final_triggering_policy
            else:
                prod_target_args["triggered_by"] = [trigger_name] if target.scheduler == SCHEDULER_LUCI else None
                prod_target_args["triggering_policy"] = final_triggering_policy

            # Both staging and prod builders are generated to retain information and make it easy to swap
            common.common_prod_builder(**prod_target_args)
            common.common_prod_builder(**staging_target_args)

def _target_platform(target):
    """Return the target platform based on the name."""
    platform = target.name.split(" ")[0]
    return platform.lower()

def _generate(repo, branch, version, testing_ref, dimensions = {}, properties = {}, triggering_policy = None, notifies = None, recipes_ref = "refs/heads/main"):
    # CQ group configurations. Only FLUTTER_RECIPES is using
    # LUCI CQ but we still need the CQ configurations for all
    # the try configurations for led recipe tests.
    if _is_default_branch(branch):
        common.cq_group(repos.GIT_REMOTE[repo])
    _recipes(repo, branch, recipes_ref)
    _builders(repo, branch, version, testing_ref, dimensions, properties, triggering_policy, notifies = notifies, recipes_ref = recipes_ref)

ci_yaml = struct(
    builder_name = _builder_name,
    ci_yaml = _ci_yaml,
    full_recipe_name = _full_recipe_name,
    generate = _generate,
    legacy_swarming_caches = _legacy_swarming_caches,
    platform_properties = _platform_properties,
    properties = _properties,
    swarming_caches = _swarming_caches,
)