# 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/.
"""
Support for running jobs that are invoked via the `run-task` script.
"""

import dataclasses
import functools
import os
from typing import Literal, Union
from typing import Optional as TOptional

from mozpack import path
from taskgraph.transforms.run.common import support_caches
from taskgraph.util.schema import (
    Schema,
    taskref_or_string_msgspec,
)
from taskgraph.util.yaml import load_yaml

from gecko_taskgraph.transforms.job import run_job_using
from gecko_taskgraph.transforms.job.common import add_tooltool, support_vcs_checkout


class RunTaskSchema(Schema, kw_only=True):
    using: Literal["run-task"]
    # Use the specified caches.
    use_caches: TOptional[Union[bool, list[str]]] = None
    # if true (the default), perform a checkout of gecko on the worker
    checkout: bool = True
    # Path to run command in. If a checkout is present, the path
    # to the checkout will be interpolated with the key `checkout`
    cwd: TOptional[str] = None
    # The sparse checkout profile to use. Value is the filename relative to
    # "sparse-profile-prefix" which defaults to "build/sparse-profiles/".
    sparse_profile: TOptional[Union[str, None]] = None
    # The relative path to the sparse profile.
    sparse_profile_prefix: TOptional[str] = None
    # Whether to use a shallow clone or not, default True (git only).
    shallow_clone: TOptional[bool] = None
    # if true, perform a checkout of a comm-central based branch inside the
    # gecko checkout
    comm_checkout: bool = False
    # The command arguments to pass to the `run-task` script, after the
    # checkout arguments.  If a list, it will be passed directly; otherwise
    # it will be included in a single argument to `bash -cx`.
    command: Union[list[taskref_or_string_msgspec], taskref_or_string_msgspec]
    # Base work directory used to set up the task.
    workdir: TOptional[str] = None
    # If not false, tooltool downloads will be enabled via relengAPIProxy
    # for either just public files, or all files. Only supported on
    # docker-worker.
    tooltool_downloads: Union[bool, Literal["public", "internal"]] = False
    # Whether to run as root. (defaults to False)
    run_as_root: TOptional[bool] = None

    def __post_init__(self):
        super().__post_init__()
        if self.tooltool_downloads is True:
            raise ValueError(
                "tooltool-downloads must be False, 'public', or 'internal'"
            )


def common_setup(config, job, taskdesc, command):
    run = job["run"]
    run_cwd = run.get("cwd")
    if run["checkout"]:
        repo_configs = config.repo_configs
        if len(repo_configs) > 1 and run["checkout"] is True:
            raise Exception("Must explicitly specify checkouts with multiple repos.")
        elif run["checkout"] is not True:
            repo_configs = {
                repo: dataclasses.replace(repo_configs[repo], **config)
                for (repo, config) in run["checkout"].items()
            }

        gecko_path = support_vcs_checkout(config, job, taskdesc, repo_configs)
        command.append(f"--gecko-checkout={gecko_path}")
        if config.params["repository_type"] == "git" and run.get("shallow-clone", True):
            command.append("--gecko-shallow-clone")

        if run_cwd:
            run_cwd = path.normpath(run_cwd.format(checkout=gecko_path))

    elif run_cwd and "{checkout}" in run_cwd:
        raise Exception(
            "Found `{{checkout}}` interpolation in `cwd` for task {name} "
            "but the task doesn't have a checkout: {cwd}".format(
                cwd=run_cwd, name=job.get("name", job.get("label"))
            )
        )

    if config.params["repository_type"] == "hg" and run["sparse-profile"]:
        sparse_profile_prefix = run.pop(
            "sparse-profile-prefix", "build/sparse-profiles"
        )
        sparse_profile_path = path.join(sparse_profile_prefix, run["sparse-profile"])
        command.append(f"--gecko-sparse-profile={sparse_profile_path}")

    if run_cwd:
        command.append(f"--task-cwd={run_cwd}")

    support_caches(config, job, taskdesc)
    taskdesc["worker"].setdefault("env", {})["MOZ_SCM_LEVEL"] = config.params["level"]


worker_defaults = {
    "checkout": True,
    "comm-checkout": False,
    "sparse-profile": None,
    "tooltool-downloads": False,
    "run-as-root": False,
}


load_yaml = functools.cache(load_yaml)


def script_url(config, script):
    if "MOZ_AUTOMATION" in os.environ and "TASK_ID" not in os.environ:
        raise Exception("TASK_ID must be defined to use run-task on generic-worker")
    task_id = os.environ.get("TASK_ID", "<TASK_ID>")
    tc_url = "http://firefox-ci-tc.services.mozilla.com"
    return f"{tc_url}/api/queue/v1/task/{task_id}/artifacts/public/{script}"


@run_job_using(
    "docker-worker", "run-task", schema=RunTaskSchema, defaults=worker_defaults
)
def docker_worker_run_task(config, job, taskdesc):
    run = job["run"]
    worker = taskdesc["worker"] = job["worker"]
    run_task_bin = (
        "run-task-git" if config.params["repository_type"] == "git" else "run-task-hg"
    )
    command = [f"/builds/worker/bin/{run_task_bin}"]
    common_setup(config, job, taskdesc, command)

    if run["tooltool-downloads"]:
        internal = run["tooltool-downloads"] == "internal"
        add_tooltool(config, job, taskdesc, internal=internal)

    run_command = run["command"]

    # dict is for the case of `{'task-reference': text_type}`.
    if isinstance(run_command, (str, dict)):
        run_command = ["bash", "-cx", run_command]
    if run["comm-checkout"]:
        command.append(
            "--comm-checkout={}/comm".format(taskdesc["worker"]["env"]["GECKO_PATH"])
        )
    if run["run-as-root"]:
        command.extend(("--user", "root", "--group", "root"))
    command.append("--")
    command.extend(run_command)
    worker["command"] = command


@run_job_using(
    "generic-worker", "run-task", schema=RunTaskSchema, defaults=worker_defaults
)
def generic_worker_run_task(config, job, taskdesc):
    run = job["run"]
    worker = taskdesc["worker"] = job["worker"]
    is_win = worker["os"] == "windows"
    is_mac = worker["os"] == "macosx"
    is_bitbar = worker["os"] == "linux-bitbar"
    is_lambda = worker["os"] == "linux-lambda"

    if run["tooltool-downloads"]:
        internal = run["tooltool-downloads"] == "internal"
        add_tooltool(config, job, taskdesc, internal=internal)

    if is_win:
        command = ["C:/mozilla-build/python3/python3.exe", "run-task"]
    elif is_mac:
        command = ["/usr/local/bin/python3", "run-task"]
    else:
        command = ["./run-task"]

    common_setup(config, job, taskdesc, command)

    worker.setdefault("mounts", [])
    run_task_bin = (
        "run-task-git" if config.params["repository_type"] == "git" else "run-task-hg"
    )
    worker["mounts"].append({
        "content": {
            "url": script_url(config, run_task_bin),
        },
        "file": "./run-task",
    })

    if (
        job.get("fetches")
        or job.get("use-uv")
        or job.get("use-python", "system") != "system"
    ):
        worker["mounts"].append({
            "content": {
                "url": script_url(config, "fetch-content"),
            },
            "file": "./fetch-content",
        })
    if run.get("checkout"):
        worker["mounts"].append({
            "content": {
                "url": script_url(config, "robustcheckout.py"),
            },
            "file": "./robustcheckout.py",
        })

    run_command = run["command"]

    # dict is for the case of `{'task-reference': text_type}`.
    if isinstance(run_command, (str, dict)):
        if is_win:
            if isinstance(run_command, dict):
                for k in run_command.keys():
                    run_command[k] = f'"{run_command[k]}"'
            else:
                run_command = f'"{run_command}"'
        run_command = ["bash", "-cx", run_command]

    if run["comm-checkout"]:
        command.append(
            "--comm-checkout={}/comm".format(taskdesc["worker"]["env"]["GECKO_PATH"])
        )

    if run["run-as-root"]:
        command.extend(("--user", "root", "--group", "root"))
    command.append("--")
    if is_bitbar:
        # Use the bitbar wrapper script which sets up the device and adb
        # environment variables
        command.append("/builds/taskcluster/script.py")
    elif is_lambda:
        command.append("/home/ltuser/taskcluster/script.py")
    command.extend(run_command)

    if is_win:
        taskref = False
        for c in command:
            if isinstance(c, dict):
                taskref = True

        if taskref:
            cmd = []
            for c in command:
                if isinstance(c, dict):
                    for v in c.values():
                        cmd.append(v)
                else:
                    cmd.append(c)
            worker["command"] = [{"artifact-reference": " ".join(cmd)}]
        else:
            worker["command"] = [" ".join(command)]
    else:
        worker["command"] = [
            ["chmod", "+x", "run-task"],
            command,
        ]