#!/usr/bin/env python

import argparse

from CRABAPI.RawCommand import crabCommand
from CRABClient.ClientExceptions import ClientException
from CRABClient.UserUtilities import setConsoleLogLevel
from httplib import HTTPException

setConsoleLogLevel(100)

p = argparse.ArgumentParser('grudgingly manage many crab tasks')
p.add_argument('--resubmit', type=bool, help='resubmit failed jobs')
p.add_argument('workdirs', nargs='*')

args = p.parse_args()

statii = []
for wd in args.workdirs:
    try:
        status = crabCommand('status', dir=wd)
        if len(status) == 0:
            print "no status for {}".format(wd)
        statii.append(status)
    except Exception:
        print "error retrieving status for {}".format(wd)
        statii.append({})

totals = [sum(s.get('jobsPerStatus', {}).values()) for s in statii]
failed = [s.get('jobsPerStatus', {}).get('failed', 0) for s in statii]
finished = [s.get('jobsPerStatus', {}).get('finished', 0) for s in statii]
running = [s.get('jobsPerStatus', {}).get('running', 0) for s in statii]

for wd, t, f, d, r in zip(args.workdirs, totals, failed, finished, running):
    try:
        print "{:30}failed:   {:5.2f} ({:5d}/{:5d})".format(wd, float(f) / t, f, t)
        print "{:30}running:  {:5.2f} ({:5d}/{:5d})".format('', float(r) / t, r, t)
        print "{:30}finished: {:5.2f} ({:5d}/{:5d})".format('', float(d) / t, d, t)
    except Exception:
        print "error processing {}".format(wd)
    print "-" * 59

ntotal = sum(totals)
nfailed = sum(failed)
nfinished = sum(finished)
nrunning = sum(running)

print "{:30}failed:   {:5.2f} ({:5d}/{:5d})".format('total', float(nfailed) / ntotal, nfailed, ntotal)
print "{:30}running:  {:5.2f} ({:5d}/{:5d})".format('', float(nrunning) / ntotal, nrunning, ntotal)
print "{:30}finished: {:5.2f} ({:5d}/{:5d})".format('', float(nfinished) / ntotal, nfinished, ntotal)

if args.resubmit:
    print
    for wd, status in zip(args.workdirs, statii):
        jobs = [str(v) for k, v in status.get('jobList', {}) if k == u'failed']
        if len(jobs) == 0:
            continue
        try:
            print "{:30}resubmitting: {}".format(wd, jobs)
            crabCommand('resubmit', jobids=','.join(jobs), dir=wd)
        except Exception:
            print "error resubmitting {}".format(wd)
