#!/usr/bin/env python
from stolos import argparse_shared as at
from stolos import zookeeper_tools as zkt, dag_tools
from stolos import log


def main(ns):
    zk = zkt.get_zkclient(ns.zookeeper_hosts)
    if ns.readd:
        zkt.readd_subtask(ns.app_name, ns.job_id, zk=zk)
    else:
        rv = zkt.maybe_add_subtask(ns.app_name, ns.job_id, zk=zk)
        if not rv:
            log.warn("Did not add subtask")


build_arg_parser = at.build_arg_parser([
    at.add_argument('--job_id', required=True),
    at.app_name(required=True, choices=dag_tools.get_task_names()),
    at.zookeeper_hosts,
    at.add_argument(
        '--readd', action='store_true', help=(
            "Try to add this to the queue if it isn't already.  Assume we've"
            " tried to add the task already")),
])


if __name__ == '__main__':
    NS = build_arg_parser().parse_args()
    main(NS)
