#!/usr/bin/env python
#
# Copyright (c) 2013, ReMake Electric ehf
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# This file implements the "malaria publish" command
"""
Publish a stream of messages and capture statistics on their timing.
"""

import argparse
import multiprocessing
import os
import random
import socket
import time

import beem.load
import beem.bridge
import beem.msgs


def my_custom_msg_generator(sequence_length):
    """
    An example of a custom msg generator.

    You must return a tuple of sequence number, topic and payload
    on each iteration.
    """
    seq = 0
    while seq < sequence_length:
        yield (seq, "magic_topic", "very boring payload")
        seq += 1


def _worker(options, proc_num, auth=None):
    """
    Wrapper to run a test and push the results back onto a queue.
    Modify this to provide custom message generation routines.
    """
    # Make a new clientid with our worker process number
    cid = "%s-%d" % (options.clientid, proc_num)
    if options.bridge:
        ts = beem.bridge.BridgingSender(options.host, options.port, cid, auth)
        # This is _probably_ what you want if you are specifying a key file
        # This would correspond with using ids as clientids, and acls
        if auth:
            cid = auth.split(":")[0]
    else:
        # FIXME - add auth support here too dummy!
        ts = beem.load.TrackingSender(options.host, options.port, cid)

    # Provide a custom generator
    #msg_gen = my_custom_msg_generator(options.msg_count)
    msg_gen = beem.msgs.createGenerator(cid, options)
    # This helps introduce jitter so you don't have many threads all in sync
    time.sleep(random.uniform(1, 10))
    ts.run(msg_gen, qos=options.qos)
    return ts.stats()


def _worker_threaded(options, proc_num, auth=None):
    ts = beem.bridge.ThreadedBridgingSender(options, proc_num, auth)
    ts.run()
    return ts.stats


def add_args(subparsers):
    parser = subparsers.add_parser(
        "publish",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description=__doc__,
        help="Publish a stream of messages")

    parser.add_argument(
        "-c", "--clientid",
        default="beem.loadr-%s-%d" % (socket.gethostname(), os.getpid()),
        help="""Set the client id of the publisher, can be useful for acls.