#!/usr/bin/env python
'''
Simulates minions via a dump file
'''

import argparse
from functools import partial
import logging
from multiprocessing import Process, cpu_count
import sys

import zmq
import zmq.eventloop.ioloop
import salt.log

from evilminions.dump_reader import DumpReader
from evilminions.evil_minion import EvilMinion

def main():
    # parse commandline switches
    parser = argparse.ArgumentParser(description='Simulates a minion given a dump file.')
    parser.add_argument('master', help='the master to connect to')
    parser.add_argument('--count', dest='count', type=int, default=10,
                       help='number of evil minions (default: 10)')
    parser.add_argument('--id-prefix', dest='prefix', default='evil',
                       help='minion id prefix for evil minions. (default: evil)')
    parser.add_argument('--id-offset', dest='offset', type=int, default=0,
                       help='minion id counter offset for evil minions. (default: 0)')
    parser.add_argument('--slowdown-factor', dest='slowdown_factor', type=float, default=0.0,
                       help='slow down evil minions to behave realistically (default is 0.0 "as fast as possible", 1.0 is "as fast as the original")')
    parser.add_argument('--processes', dest='processes', type=int, default=cpu_count(),
                       help='number of concurrent processes (default is the CPU count: %d)' % cpu_count())
    parser.add_argument('--keysize', dest='keysize', type=int, default=2048,
                       help='size of Salt keys generated for each evil minions (default: 2048)')
    parser.add_argument('--dump-path', dest='dump_path', default='/tmp/minion-dump.mp',
                       help='path for the dump file (default: /tmp/minion-dump.mp)')
    args = parser.parse_args()

    # set up logging
    salt.log.setup_console_logger(log_level='debug')

    # set up processes
    for chunk in minion_chunks(args.count, args.processes):
        p = Process(target=process_main, args=(args, chunk))
        p.start()

def process_main(args, chunk):
    '''Per-process entry point'''
    # set up logging
    log = logging.getLogger(__name__)
    log.debug("Starting pool: %s" % chunk)

    # set up the IO loop
    zmq.eventloop.ioloop.install()
    io_loop = zmq.eventloop.ioloop.ZMQIOLoop.current()

    # HACK (brutal): force exactly one ZeroMQ context for the whole process
    # this avoids the creation of 4 threads per evil minion
    context = zmq.Context()
    context.term = lambda: None
    zmq.Context = lambda: context

    # set up the dump
    dump_reader = DumpReader(args.dump_path)

    # start the evil!
    for i in chunk:
        evil_minion = EvilMinion(args.master,
                                 '{}-{}'.format(args.prefix, i + args.offset),
                                 args.slowdown_factor,
                                 dump_reader,
                                 io_loop,
                                 keysize=args.keysize)
        io_loop.spawn_callback(evil_minion.start)
    io_loop.start()

def minion_chunks(count, processes):
    '''Returns chunks of minion indexes, per process (eg. 3 CPUs and 4 processes: [[0, 1], [2], [3]])'''
    minions_per_process = count / processes
    rest = count % processes
    lengths = [minions_per_process + 1] * rest + [minions_per_process] * (processes - rest)
    starts = [sum(lengths[:i]) for i in range(processes)]

    indexes = list(range(count))

    return [indexes[starts[i]:starts[i] + lengths[i]] for i in range(processes)]

if __name__ == "__main__":
    main()
