# 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/.
"""
Source-test jobs can run on multiple platforms.  These transforms allow jobs
with either `platform` or a list of `platforms`, and set the appropriate
treeherder configuration and attributes for that platform.
"""


import copy
import os

from taskgraph.transforms.base import TransformSequence
from taskgraph.util.attributes import keymatch
from taskgraph.util.schema import Schema, optionally_keyed_by, resolve_keyed_by
from taskgraph.util.treeherder import join_symbol, split_symbol
from voluptuous import Any, Extra, Optional, Required

from gecko_taskgraph.transforms.job import job_description_schema

source_test_description_schema = Schema(
    {
        # most fields are passed directly through as job fields, and are not
        # repeated here
        Extra: object,
        # The platform on which this task runs.  This will be used to set up attributes
        # (for try selection) and treeherder metadata (for display).  If given as a list,
        # the job will be "split" into multiple tasks, one with each platform.
        Required("platform"): Any(str, [str]),
        # Build labels required for the task. If this key is provided it must
        # contain a build label for the task platform.
        # The task will then depend on a build task, and the installer url will be
        # saved to the GECKO_INSTALLER_URL environment variable.
        Optional("require-build"): optionally_keyed_by("project", {str: str}),
        # These fields can be keyed by "platform", and are otherwise identical to
        # job descriptions.
        Required("worker-type"): optionally_keyed_by(
            "platform", job_description_schema["worker-type"]
        ),
        Required("worker"): optionally_keyed_by(
            "platform", job_description_schema["worker"]
        ),
        Optional("python-version"): [int],
        Optional("dependencies"): {
            k: optionally_keyed_by("platform", v)
            for k, v in job_description_schema["dependencies"].items()
        },
        # A list of artifacts to install from 'fetch' tasks.
        Optional("fetches"): {
            str: optionally_keyed_by(
                "platform", job_description_schema["fetches"][str]
            ),
        },
    }
)

transforms = TransformSequence()

transforms.add_validate(source_test_description_schema)


@transforms.add
def set_job_name(config, jobs):
    for job in jobs:
        if "job-from" in job and job["job-from"] != "kind.yml":
            from_name = os.path.splitext(job["job-from"])[0]
            job["name"] = "{}-{}".format(from_name, job["name"])
        yield job


@transforms.add
def expand_platforms(config, jobs):
    for job in jobs:
        if isinstance(job["platform"], str):
            yield job
            continue

        for platform in job["platform"]:
            pjob = copy.deepcopy(job)
            pjob["platform"] = platform

            if "name" in pjob:
                pjob["name"] = "{}-{}".format(pjob["name"], platform)
            else:
                pjob["label"] = "{}-{}".format(pjob["label"], platform)
            yield pjob


@transforms.add
def split_python(config, jobs):
    for job in jobs:
        key = "python-version"
        versions = job.pop(key, [])
        if not versions:
            yield job
            continue
        for version in versions:
            group = f"py{version}"
            pyjob = copy.deepcopy(job)
            if "name" in pyjob:
                pyjob["name"] += f"-{group}"
            else:
                pyjob["label"] += f"-{group}"
            symbol = split_symbol(pyjob["treeherder"]["symbol"])[1]
            pyjob["treeherder"]["symbol"] = join_symbol(group, symbol)
            pyjob["run"][key] = version
            yield pyjob


@transforms.add
def split_jsshell(config, jobs):
    all_shells = {"sm": "Spidermonkey", "v8": "Google V8"}

    for job in jobs:
        if not job["name"].startswith("jsshell"):
            yield job
            continue

        test = job.pop("test")
        for shell in job.get("shell", all_shells.keys()):
            assert shell in all_shells

            new_job = copy.deepcopy(job)
            new_job["name"] = "{}-{}".format(new_job["name"], shell)
            new_job["description"] = "{} on {}".format(
                new_job["description"], all_shells[shell]
            )
            new_job["shell"] = shell

            group = f"js-bench-{shell}"
            symbol = split_symbol(new_job["treeherder"]["symbol"])[1]
            new_job["treeherder"]["symbol"] = join_symbol(group, symbol)

            run = new_job["run"]
            run["mach"] = run["mach"].format(
                shell=shell, SHELL=shell.upper(), test=test
            )
            yield new_job


def add_build_dependency(config, job):
    """
    Add build dependency to the job and installer_url to env.
    """
    key = job["platform"]
    build_labels = job.pop("require-build", {})
    matches = keymatch(build_labels, key)
    if not matches:
        raise Exception(
            "No build platform found. "
            "Define 'require-build' for {} in the task config.".format(key)
        )

    if len(matches) > 1:
        raise Exception(f"More than one build platform found for '{key}'.")

    label = matches[0]
    deps = job.setdefault("dependencies", {})
    deps.update({"build": label})


@transforms.add
def handle_platform(config, jobs):
    """
    Handle the 'platform' property, setting up treeherder context as well as
    try-related attributes.
    """
    fields = [
        "always-target",
        "fetches.toolchain",
        "require-build",
        "worker-type",
        "worker",
    ]

    for job in jobs:
        platform = job["platform"]

        for field in fields:
            resolve_keyed_by(
                job, field, item_name=job["name"], project=config.params["project"]
            )
        for field in job.get("dependencies", {}):
            resolve_keyed_by(
                job,
                f"dependencies.{field}",
                item_name=job["name"],
                project=config.params["project"],
            )

        if "treeherder" in job:
            job["treeherder"].setdefault("platform", platform)

        if "require-build" in job:
            add_build_dependency(config, job)

        del job["platform"]
        yield job


@transforms.add
def handle_shell(config, jobs):
    """
    Handle the 'shell' property.
    """
    fields = [
        "run-on-projects",
        "worker.env",
    ]

    for job in jobs:
        if not job.get("shell"):
            yield job
            continue

        for field in fields:
            resolve_keyed_by(job, field, item_name=job["name"])

        del job["shell"]
        yield job


@transforms.add
def set_code_review_env(config, jobs):
    """
    Add a CODE_REVIEW environment variable when running in code-review bot mode
    """
    is_code_review = config.params["target_tasks_method"] == "codereview"

    for job in jobs:
        attrs = job.get("attributes", {})
        if is_code_review and attrs.get("code-review") is True:
            env = job["worker"].setdefault("env", {})
            env["CODE_REVIEW"] = "1"

        yield job


@transforms.add
def set_worker_exit_code(config, jobs):
    for job in jobs:
        worker = job["worker"]
        worker.setdefault("retry-exit-status", [])
        if 137 not in worker["retry-exit-status"]:
            worker["retry-exit-status"].append(137)
        yield job


@transforms.add
def remove_optimization_on_central(config, jobs):
    """
    For pushes to mozilla-central run all source-test tasks that are enabled for
    code-review in order to have the code-review bot populate the DB according
    with the push hash.
    """
    if (
        config.params["project"] != "mozilla-central"
        or config.params["tasks_for"] != "hg-push"
    ):
        yield from jobs
        return

    for job in jobs:
        if not job.get("attributes", {}).get("code-review", False):
            yield job
            continue
        if "when" not in job:
            yield job
            continue
        del job["when"]
        yield job