#!/usr/bin/env python
import foopy_fabric
from fabric.api import env
from fabric.context_managers import settings
from fabric.colors import red
from Crypto.Random import atfork

FAIL = red('[FAIL]')

_devices_map = dict()
_foopy_map = dict()


def _calc_maps(devices_file):
    global _devices_map, _foopy_map
    all_devices = json.load(urllib.urlopen(devices_file))
    _devices_map = dict([[x, all_devices[x]] for x in all_devices
                         if 'foopy' in all_devices[x]
                         and all_devices[x]['foopy'] is not None])
    # Extract foopies from devices
    all_foopies = [all_devices[x]['foopy'] for x in all_devices
                   if 'foopy' in all_devices[x]]
    # Use set to trim dupes
    all_foopies = set(all_foopies) - set(["None"])
    # For sanity when using -H all, we revert to list and sort
    all_foopies = [x for x in all_foopies]
    all_foopies.sort()
    for foopy in all_foopies:
        _foopy_map[foopy] = dict([[x, _devices_map[x]] for x in _devices_map
                                  if _devices_map[x]['foopy'] == foopy])


def _calc_selected(hosts, devices):
    global _foopy_map, _devices_map
    from copy import deepcopy
    selected = dict()
    if 'all' in hosts or 'all' in devices:
        selected = deepcopy(_foopy_map)
        return selected

    for f in hosts:
        if f in _foopy_map:
            selected[f] = deepcopy(_foopy_map[f])
    for d in devices:
        if d not in _devices_map:
            sys.stderr.write("Warning: device '%s' not found in device map\n" % d)
            continue
        if 'foopy' not in _devices_map[d]:
            continue
        if _devices_map[d]['foopy'] not in _foopy_map:
            continue
        if _devices_map[d]['foopy'] in selected:
            selected[_devices_map[d]['foopy']][d] = deepcopy(d)
        else:  # No mapping for this foopy
            selected[_devices_map[d]['foopy']] = dict()
            selected[_devices_map[d]['foopy']][d] = deepcopy(d)
    return selected


def run_action_on_foopy(action, foopy):
    atfork()
    try:
        action_func = getattr(foopy_fabric, action)
        with settings(host_string="%s.build.mozilla.org" % foopy):
            if options.argument:
                action_func(foopy, options.argument)
            else:
                action_func(foopy)
            return True
    except AttributeError:
        print FAIL, "[%s] %s action is not defined." % (foopy, action)
        return False
    except:
        import traceback
        print "Failed to run", action, "on", foopy
        print traceback.format_exc()
        return False


def run_action_on_devices(action, foopy_dict):
    atfork()
    action_func = getattr(foopy_fabric, action)
    try:
        with settings(host_string="%s.build.mozilla.org" % foopy_dict['host']):
            for device in foopy_dict['devices']:
                if hasattr(action_func, "needs_device_dict"):
                    action_func(device, _devices_map[device])
                else:
                    action_func(device)
            return True
        return False
    except AttributeError:
        print FAIL, "[%s] %s action is not defined." % (foopy_dict['host'], action)
        return False
    except:
        import traceback
        print "Failed to run", action, "on", foopy_dict['host']
        print traceback.format_exc()
        return False

if __name__ == '__main__':
    from optparse import OptionParser
    import sys
    import urllib
    try:
        import simplejson as json
    except:
        import json

    parser = OptionParser("""%%prog [options] action [action ...]""")

    parser.set_defaults(
        hosts=[],
        devices=[],
        concurrency=1,
    )
    parser.add_option("-f", "--devices-file", dest="devices_file", help="list/url of devices.json")
    parser.add_option("-H", "--host", dest="hosts", action="append")
    parser.add_option("-D", "--device", dest="devices", action="append")
    parser.add_option("-j", dest="concurrency", type="int")
    parser.add_option("-a", dest="argument", action="store")

    options, actions = parser.parse_args()

    if options.concurrency > 1:
        import multiprocessing

    if not options.devices_file:
        parser.error("devices-file is required")

    if not actions:
        parser.error("at least one action is required")

    _calc_maps(options.devices_file)

    selected = _calc_selected(options.hosts, options.devices)

    if len(selected.keys()) == 0:
        parser.error("You need to specify a foopy via -H or a device via -D")

    env.user = 'cltbld'

    for action in actions:
        fn = run_action_on_foopy
        genArg = lambda x: x
        if hasattr(getattr(foopy_fabric, action), "per_device"):
            fn = run_action_on_devices
            genArg = lambda x: {'host': x, 'devices': selected[x]}

        if options.concurrency == 1:
            for foopy in sorted(selected.keys()):
                fn(action, genArg(foopy))
        else:
            # Don't prompt for passwords when forking
            env.abort_on_prompts = True

            p = multiprocessing.Pool(processes=options.concurrency)
            results = []
            for foopy in sorted(selected.keys()):
                result = p.apply_async(fn, (action, genArg(foopy)))
                results.append((foopy, result))
            p.close()
            failed = False
            for foopy, result in results:
                if not result.get():
                    print foopy, "FAILED"
                    failed = True
            p.join()
            if failed:
                sys.exit(1)