# -*- Mode: python; indent-tabs-mode: nil; tab-width: 40 -*-
# vim: set filetype=python:
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

option(
    env="MOZ_FETCHES_DIR",
    nargs=1,
    when="MOZ_AUTOMATION",
    help="Directory containing fetched artifacts",
)


@depends("MOZ_FETCHES_DIR", when="MOZ_AUTOMATION")
def moz_fetches_dir(value):
    if value:
        return value[0]


@depends(vcs_checkout_type, milestone.is_nightly, "MOZ_AUTOMATION")
def bootstrap_default(vcs_checkout_type, is_nightly, automation):
    if automation:
        return False
    # We only enable if building off a VCS checkout of central.
    if is_nightly and vcs_checkout_type:
        return True


option(
    "--enable-bootstrap",
    nargs="*",
    default=bootstrap_default,
    help="{Automatically bootstrap or update some toolchains|Disable bootstrap or update of toolchains}",
)


@depends_if("--enable-bootstrap")
def enable_bootstrap(bootstrap):
    include = set()
    exclude = set()
    for item in bootstrap:
        if item.startswith("-"):
            exclude.add(item.lstrip("-"))
        else:
            include.add(item)

    def match(name):
        if name in exclude:
            return False
        if include and name in include:
            return True
        return not bool(include)

    return match


@depends(developer_options, "--enable-bootstrap", moz_fetches_dir)
def bootstrap_search_path_order(developer_options, bootstrap, moz_fetches_dir):
    if moz_fetches_dir:
        log.debug("Prioritizing MOZ_FETCHES_DIR in toolchain path.")
        return "prepend"

    if bootstrap:
        log.debug(
            "Prioritizing mozbuild state dir in toolchain paths because "
            "bootstrap mode is enabled."
        )
        return "maybe-prepend"

    if developer_options:
        log.debug(
            "Prioritizing mozbuild state dir in toolchain paths because "
            "you are not building in release mode."
        )
        return "prepend"

    log.debug(
        "Prioritizing system over mozbuild state dir in "
        "toolchain paths because you are building in "
        "release mode."
    )
    return "append"


toolchains_base_dir = moz_fetches_dir | mozbuild_state_path


@dependable
@imports("os")
@imports(_from="os", _import="environ")
def original_path():
    return environ["PATH"].split(os.pathsep)


@depends(host, when="--enable-bootstrap")
@imports("os")
@imports("traceback")
@imports(_from="mozbuild.toolchains", _import="toolchain_task_definitions")
@imports(_from="__builtin__", _import="Exception")
def bootstrap_toolchain_tasks(host):
    prefix = {
        ("x86_64", "GNU", "Linux"): "linux64",
        ("x86_64", "OSX", "Darwin"): "macosx64",
        ("aarch64", "OSX", "Darwin"): "macosx64-aarch64",
        ("x86_64", "WINNT", "WINNT"): "win64",
        ("aarch64", "WINNT", "WINNT"): "win64-aarch64",
    }.get((host.cpu, host.os, host.kernel))
    try:
        tasks = toolchain_task_definitions()
    except Exception as e:
        message = traceback.format_exc()
        log.warning(str(e))
        log.debug(message)
        return None

    def task_data(t):
        result = {
            "index": t.optimization["index-search"],
            "artifact": t.attributes["toolchain-artifact"],
        }
        command = t.attributes.get("toolchain-command")
        if command:
            result["command"] = command
        return result

    # We only want to use toolchains annotated with "local-toolchain". We also limit the
    # amount of data to what we use, so that trace logs can be more useful.
    tasks = {
        k: task_data(t)
        for k, t in tasks.items()
        if t.attributes.get("local-toolchain") and "index-search" in t.optimization
    }

    return namespace(prefix=prefix, tasks=tasks)


@template
def bootstrap_path(path, **kwargs):
    when = kwargs.pop("when", None)
    if kwargs:
        configure_error("bootstrap_path only takes `when` as a keyword argument")

    @depends(
        enable_bootstrap,
        toolchains_base_dir,
        moz_fetches_dir,
        bootstrap_toolchain_tasks,
        build_environment,
        dependable(path),
        when=when,
    )
    @imports("os")
    @imports("subprocess")
    @imports("sys")
    @imports(_from="mozbuild.util", _import="ensureParentDir")
    @imports(_from="importlib", _import="import_module")
    @imports(_from="shutil", _import="rmtree")
    @imports(_from="__builtin__", _import="open")
    @imports(_from="__builtin__", _import="Exception")
    def bootstrap_path(
        bootstrap, toolchains_base_dir, moz_fetches_dir, tasks, build_env, path
    ):
        if not path:
            return
        path_parts = path.split("/")
        path_prefix = ""
        # Small hack until clang-tidy stops being a separate toolchain in a
        # weird location.
        if path_parts[0] == "clang-tools":
            path_prefix = path_parts.pop(0)

        def try_bootstrap(exists):
            if not tasks:
                return False
            prefixes = [""]
            if tasks.prefix:
                prefixes.insert(0, "{}-".format(tasks.prefix))
            for prefix in prefixes:
                label = "toolchain-{}{}".format(prefix, path_parts[0])
                task = tasks.tasks.get(label)
                if task:
                    break
            log.debug("Trying to bootstrap %s", label)
            if not task:
                return False
            task_index = task["index"]
            log.debug("Resolved %s to %s", label, task_index[0])
            task_index = task_index[0].split(".")[-1]
            artifact = task["artifact"]
            # `mach artifact toolchain` doesn't support authentication for
            # private artifacts. Some toolchains may provide a command that can be
            # used for local production of the artifact.
            command = None
            if not artifact.startswith("public/"):
                command = task.get("command")
                if not command:
                    log.debug("Cannot bootstrap %s: not a public artifact", label)
                    return False
            index_file = os.path.join(toolchains_base_dir, "indices", path_parts[0])
            try:
                with open(index_file) as fh:
                    index = fh.read().strip()
            except Exception:
                # On automation, if there's an artifact in MOZ_FETCHES_DIR, we assume it's
                # up-to-date.
                index = task_index if moz_fetches_dir else None
            if index == task_index and exists:
                log.debug("%s is up-to-date", label)
                return True
            # Manually import with import_module so that we can gracefully disable bootstrap
            # when e.g. building from a js standalone tarball, that doesn't contain the
            # taskgraph code. In those cases, `mach artifact toolchain --from-build` would
            # also fail.
            task_id = None
            if not command:
                try:
                    IndexSearch = import_module(
                        "gecko_taskgraph.optimize.strategies"
                    ).IndexSearch
                except Exception:
                    log.debug("Cannot bootstrap %s: missing taskgraph module", label)
                    return False
                task_id = IndexSearch().should_replace_task(
                    task, {}, None, task["index"]
                )
            if task_id:
                # If we found the task in the index, use the `mach artifact toolchain`
                # fast path.
                command = [
                    "artifact",
                    "toolchain",
                    "--from-task",
                    f"{task_id}:{artifact}",
                ]
            elif command:
                # For private local toolchains, run the associated command.
                command = (
                    [
                        "python",
                        "--virtualenv",
                        "build",
                        os.path.join(
                            build_env.topsrcdir,
                            "taskcluster/scripts/misc",
                            command["script"],
                        ),
                    ]
                    + command["arguments"]
                    + [path_parts[0]]
                )
                # Clean up anything that was bootstrapped previously before going
                # forward. In other cases, that's taken care of by mach artifact toolchain.
                rmtree(
                    os.path.join(toolchains_base_dir, path_prefix, path_parts[0]),
                    ignore_errors=True,
                )
            else:
                # Otherwise, use the slower path, which will print a better error than
                # we would be able to.
                command = ["artifact", "toolchain", "--from-build", label]

            log.info(
                "%s bootstrapped toolchain in %s",
                "Updating" if exists else "Installing",
                os.path.join(toolchains_base_dir, path_prefix, path_parts[0]),
            )
            os.makedirs(os.path.join(toolchains_base_dir, path_prefix), exist_ok=True)
            subprocess.run(
                [
                    sys.executable,
                    os.path.join(build_env.topsrcdir, "mach"),
                    "--log-no-times",
                ]
                + command,
                cwd=os.path.join(toolchains_base_dir, path_prefix),
                check=True,
            )
            ensureParentDir(index_file)
            with open(index_file, "w") as fh:
                fh.write(task_index)
            return True

        path = os.path.join(toolchains_base_dir, path_prefix, *path_parts)
        if bootstrap and bootstrap(path_parts[0]):
            try:
                if not try_bootstrap(os.path.exists(path)):
                    # If there aren't toolchain artifacts to use for this build,
                    # don't return a path.
                    return None
            except Exception as e:
                log.error("%s", e)
                die("If you can't fix the above, retry with --disable-bootstrap.")
        # We re-test whether the path exists because it may have been created by
        # try_bootstrap. Automation will not have gone through the bootstrap
        # process, but we want to return the path if it exists.
        if os.path.exists(path):
            return path

    return bootstrap_path


@template
def bootstrap_search_path(path, paths=original_path, **kwargs):
    @depends(
        enable_bootstrap,
        bootstrap_path(path, **kwargs),
        bootstrap_search_path_order,
        paths,
        original_path,
    )
    def bootstrap_search_path(bootstrap, path, order, paths, original_path):
        if paths is None:
            paths = original_path
        if not path:
            return paths
        if order == "maybe-prepend":
            if bootstrap(path.split("/")[0]):
                order = "prepend"
            else:
                order = "append"
        if order == "prepend":
            return [path] + paths
        return paths + [path]

    return bootstrap_search_path