import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401


def check_context():
    """
    Check if GetFailedTasks is in the context, via checking if ${GetFailedTasks} exists, else fail with message.

    Returns:
        A list of the failed tasks generated by GetFailedTasks command.

    """
    incidents = demisto.incidents()
    if not incidents:
        raise DemistoException("No incidents were found. Make sure you are running this task from an existing incident.")
    incident_id = incidents[0]["id"]
    failed_tasks = demisto.executeCommand("getContext", {"id": incident_id})
    if is_error(failed_tasks):
        raise DemistoException(f"Error while retrieving context data. Error:\n{get_error(failed_tasks)}")
    failed_tasks = failed_tasks[0].get("Contents", {}).get("context", {}).get("GetFailedTasks")
    if not failed_tasks:
        raise DemistoException(
            "Couldn't find failed tasks in the context under the key GetFailedTasks. Please run !GetFailedTasks and try again."
        )
    return failed_tasks


def remove_exclusion(failed_tasks: list, playbook_exclusion: list):
    """
    Checks if one of the failed tasks is from an excluded playbook and if so removes it from the list.

    Args:
        failed_tasks: A list of failed tasks.
        playbook_exclusion: A list of names of playbooks to exclude.

    Returns:
        The modified list.
    """

    for playbook in playbook_exclusion:
        for task in failed_tasks:
            if playbook in task["Playbook Name"]:
                failed_tasks.remove(task)
    return failed_tasks


def restart_tasks(failed_tasks: list, sleep_time: int, group_size: int):
    """
    Reopen the given tasks and re-run them. The functions sleeps for the given time if group size has been hit.

    Args:
        failed_tasks: The list of failed tasks.
        sleep_time: The amount of seconds to sleep after restarting the group of tasks.
        group_size: The group size to be reached before sleeping .

    Returns:
        The number of failed tasks that were reopened and a dict with the data of these tasks.
    """
    restarted_tasks_count = 0
    restarted_tasks = []
    is_xsoar_version_6_2 = is_demisto_version_ge("6.2")
    for task in failed_tasks:
        task_id, incident_id, playbook_name, task_name = (
            task["Task ID"],
            task["Incident ID"],
            task["Playbook Name"],
            task["Task Name"],
        )
        demisto.info(f"Restarting task with id: {task_id} and incident id: {incident_id}")
        demisto.executeCommand("taskReopen", {"id": task_id, "incidentId": incident_id})
        body = {"invId": incident_id, "inTaskID": task_id}

        if is_xsoar_version_6_2:
            body = {"taskinfo": body}

        demisto.executeCommand("core-api-post", {"uri": "inv-playbook/task/execute", "body": json.dumps(body)})
        restarted_tasks.append(
            {"IncidentID": incident_id, "TaskID": task_id, "PlaybookName": playbook_name, "TaskName": task_name}
        )
        restarted_tasks_count += 1

        # Sleep if the group size has been hit
        if restarted_tasks_count % group_size == 0:
            demisto.info("Sleeping")
            time.sleep(sleep_time)  # pylint: disable=E9003

    return restarted_tasks_count, restarted_tasks


def main():
    args = demisto.args()

    # Get Arguments
    playbook_exclusion = argToList(args.get("playbook_exclusion"))
    sleep_time = int(args.get("sleep_time"))
    incident_limit = int(args.get("incident_limit"))
    group_size = int(args.get("group_size"))

    try:
        if group_size == 0:
            raise DemistoException("The group size argument should be 1 or higher.")

        # Get Context for Failed Tasks
        failed_tasks = check_context()
        # Remove Excluded Playbooks And Limit
        failed_tasks = remove_exclusion(failed_tasks, playbook_exclusion)[:incident_limit]

        # Restart the tasks, make sure the number of incidents does not exceed the limit
        restarted_tasks_count, restarted_tasks = restart_tasks(failed_tasks, sleep_time, group_size)
        human_readable = tableToMarkdown(
            "Tasks Restarted",
            restarted_tasks,
            headers=["IncidentID", "PlaybookName", "TaskName", "TaskID"],
            headerTransform=pascalToSpace,
        )

        return_results(
            CommandResults(
                readable_output=human_readable,
                outputs_prefix="RestartedTasks",
                outputs={"Total": restarted_tasks_count, "Task": restarted_tasks},
            )
        )
    except DemistoException as e:
        return_error(f"Failed while trying to restart failed tasks. Error: {e}", error=traceback.format_exc())


if __name__ in ("__main__", "__builtin__", "builtins"):
    main()