#!/usr/bin/env python3
# Copyright 2019 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Updates the status of a tryjob."""

import argparse
import enum
import json
import os
import subprocess
import sys

import chroot
import test_helpers


class TryjobStatus(enum.Enum):
    """Values for the 'status' field of a tryjob."""

    GOOD = "good"
    BAD = "bad"
    PENDING = "pending"
    SKIP = "skip"

    # Executes the script passed into the command line (this script's exit code
    # determines the 'status' value of the tryjob).
    CUSTOM_SCRIPT = "custom_script"


class CustomScriptStatus(enum.Enum):
    """Exit code values of a custom script."""

    # NOTE: Not using 1 for 'bad' because the custom script can raise an
    # exception which would cause the exit code of the script to be 1, so the
    # tryjob's 'status' would be updated when there is an exception.
    #
    # Exit codes are as follows:
    #   0: 'good'
    #   124: 'bad'
    #   125: 'skip'
    GOOD = 0
    BAD = 124
    SKIP = 125


custom_script_exit_value_mapping = {
    CustomScriptStatus.GOOD.value: TryjobStatus.GOOD.value,
    CustomScriptStatus.BAD.value: TryjobStatus.BAD.value,
    CustomScriptStatus.SKIP.value: TryjobStatus.SKIP.value,
}


def GetCommandLineArgs():
    """Parses the command line for the command line arguments."""

    # Default absoute path to the chroot if not specified.
    cros_root = os.path.expanduser("~")
    cros_root = os.path.join(cros_root, "chromiumos")

    # Create parser and add optional command-line arguments.
    parser = argparse.ArgumentParser(
        description="Updates the status of a tryjob."
    )

    # Add argument for the JSON file to use for the update of a tryjob.
    parser.add_argument(
        "--status_file",
        required=True,
        help="The absolute path to the JSON file that contains the tryjobs "
        "used for bisecting LLVM.",
    )

    # Add argument that sets the 'status' field to that value.
    parser.add_argument(
        "--set_status",
        required=True,
        choices=[tryjob_status.value for tryjob_status in TryjobStatus],
        help='Sets the "status" field of the tryjob.',
    )

    # Add argument that determines which revision to search for in the list of
    # tryjobs.
    parser.add_argument(
        "--revision",
        required=True,
        type=int,
        help="The revision to set its status.",
    )

    # Add argument for the custom script to execute for the 'custom_script'
    # option in '--set_status'.
    parser.add_argument(
        "--custom_script",
        help="The absolute path to the custom script to execute (its exit code "
        'should be %d for "good", %d for "bad", or %d for "skip")'
        % (
            CustomScriptStatus.GOOD.value,
            CustomScriptStatus.BAD.value,
            CustomScriptStatus.SKIP.value,
        ),
    )

    args_output = parser.parse_args()

    if not (
        os.path.isfile(
            args_output.status_file
            and not args_output.status_file.endswith(".json")
        )
    ):
        raise ValueError(
            'File does not exist or does not ending in ".json" '
            ": %s" % args_output.status_file
        )

    if (
        args_output.set_status == TryjobStatus.CUSTOM_SCRIPT.value
        and not args_output.custom_script
    ):
        raise ValueError(
            "Please provide the absolute path to the script to " "execute."
        )

    return args_output


def FindTryjobIndex(revision, tryjobs_list):
    """Searches the list of tryjob dictionaries to find 'revision'.

    Uses the key 'rev' for each dictionary and compares the value against
    'revision.'

    Args:
        revision: The revision to search for in the tryjobs.
        tryjobs_list: A list of tryjob dictionaries of the format:
        {
            'rev' : [REVISION],
            'url' : [URL_OF_CL],
            'cl' : [CL_NUMBER],
            'link' : [TRYJOB_LINK],
            'status' : [TRYJOB_STATUS],
            'buildbucket_id': [BUILDBUCKET_ID]
        }

    Returns:
        The index within the list or None to indicate it was not found.
    """

    for cur_index, cur_tryjob_dict in enumerate(tryjobs_list):
        if cur_tryjob_dict["rev"] == revision:
            return cur_index

    return None


def GetCustomScriptResult(custom_script, status_file, tryjob_contents):
    """Returns the conversion of the exit code of the custom script.

    Args:
        custom_script: Absolute path to the script to be executed.
        status_file: Absolute path to the file that contains information about
        the bisection of LLVM.
        tryjob_contents: A dictionary of the contents of the tryjob (e.g.
        'status', 'url', 'link', 'buildbucket_id', etc.).

    Returns:
        The exit code conversion to either return 'good', 'bad', or 'skip'.

    Raises:
        ValueError: The custom script failed to provide the correct exit code.
    """

    # Create a temporary file to write the contents of the tryjob at index
    # 'tryjob_index' (the temporary file path will be passed into the custom
    # script as a command line argument).
    with test_helpers.CreateTemporaryJsonFile() as temp_json_file:
        with open(temp_json_file, "w", encoding="utf-8") as tryjob_file:
            json.dump(
                tryjob_contents, tryjob_file, indent=4, separators=(",", ": ")
            )

        exec_script_cmd = [custom_script, temp_json_file]

        # Execute the custom script to get the exit code.
        with subprocess.Popen(
            exec_script_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        ) as exec_script_cmd_obj:
            _, stderr = exec_script_cmd_obj.communicate()

        # Invalid exit code by the custom script.
        if (
            exec_script_cmd_obj.returncode
            not in custom_script_exit_value_mapping
        ):
            # Save the .JSON file to the directory of 'status_file'.
            name_of_json_file = os.path.join(
                os.path.dirname(status_file), os.path.basename(temp_json_file)
            )

            os.rename(temp_json_file, name_of_json_file)

            raise ValueError(
                "Custom script %s exit code %d did not match "
                'any of the expected exit codes: %d for "good", %d '
                'for "bad", or %d for "skip".\nPlease check %s for information '
                "about the tryjob: %s"
                % (
                    custom_script,
                    exec_script_cmd_obj.returncode,
                    CustomScriptStatus.GOOD.value,
                    CustomScriptStatus.BAD.value,
                    CustomScriptStatus.SKIP.value,
                    name_of_json_file,
                    stderr,
                )
            )

    return custom_script_exit_value_mapping[exec_script_cmd_obj.returncode]


def UpdateTryjobStatus(revision, set_status, status_file, custom_script):
    """Updates a tryjob's 'status' field based off of 'set_status'.

    Args:
        revision: The revision associated with the tryjob.
        set_status: What to update the 'status' field to.
            Ex: TryjobStatus.Good, TryjobStatus.BAD, TryjobStatus.PENDING, or
            TryjobStatus.
        status_file: The .JSON file that contains the tryjobs.
        custom_script: The absolute path to a script that will be executed
        which will determine the 'status' value of the tryjob.
    """

    # Format of 'bisect_contents':
    # {
    #   'start': [START_REVISION_OF_BISECTION]
    #   'end': [END_REVISION_OF_BISECTION]
    #   'jobs' : [
    #       {[TRYJOB_INFORMATION]},
    #       {[TRYJOB_INFORMATION]},
    #       ...,
    #       {[TRYJOB_INFORMATION]}
    #   ]
    # }
    with open(status_file, encoding="utf-8") as tryjobs:
        bisect_contents = json.load(tryjobs)

    if not bisect_contents["jobs"]:
        sys.exit("No tryjobs in %s" % status_file)

    tryjob_index = FindTryjobIndex(revision, bisect_contents["jobs"])

    # 'FindTryjobIndex()' returns None if the revision was not found.
    if tryjob_index is None:
        raise ValueError(
            "Unable to find tryjob for %d in %s" % (revision, status_file)
        )

    # Set 'status' depending on 'set_status' for the tryjob.
    if set_status == TryjobStatus.GOOD:
        bisect_contents["jobs"][tryjob_index][
            "status"
        ] = TryjobStatus.GOOD.value
    elif set_status == TryjobStatus.BAD:
        bisect_contents["jobs"][tryjob_index]["status"] = TryjobStatus.BAD.value
    elif set_status == TryjobStatus.PENDING:
        bisect_contents["jobs"][tryjob_index][
            "status"
        ] = TryjobStatus.PENDING.value
    elif set_status == TryjobStatus.SKIP:
        bisect_contents["jobs"][tryjob_index][
            "status"
        ] = TryjobStatus.SKIP.value
    elif set_status == TryjobStatus.CUSTOM_SCRIPT:
        bisect_contents["jobs"][tryjob_index]["status"] = GetCustomScriptResult(
            custom_script, status_file, bisect_contents["jobs"][tryjob_index]
        )
    else:
        raise ValueError(
            'Invalid "set_status" option provided: %s' % set_status
        )

    with open(status_file, "w", encoding="utf-8") as update_tryjobs:
        json.dump(
            bisect_contents, update_tryjobs, indent=4, separators=(",", ": ")
        )


def main():
    """Updates the status of a tryjob."""

    chroot.VerifyOutsideChroot()

    args_output = GetCommandLineArgs()

    UpdateTryjobStatus(
        args_output.revision,
        TryjobStatus(args_output.set_status),
        args_output.status_file,
        args_output.custom_script,
    )


if __name__ == "__main__":
    main()