import txamqp
from twisted.internet.protocol import ClientCreator
from txamqp.protocol import AMQClient
from twisted.internet import defer, reactor
from txamqp.content import Content
from txamqp.client import TwistedDelegate
import time


messages = 10000
start = time.time()

def receive(msg, queue):
    print('receiving: %s' % msg.content.body)
    if msg.content.body == str(messages):
        stop = time.time()
        total = stop - start
        mps = messages/total
        print("total time (%s messages): %s seconds" % (messages, total))
        print("messages per second: %s" % mps)
        reactor.stop()
    else:
        consume_queue(queue)

def consume_queue(queue):
    queue.get().addCallback(receive, queue)

def send(channel, content):
    reactor.callFromThread(channel.basic_publish, routing_key='queue', content=content)
    print('sending: %s' % content.body)

def send_msgs_from_for_loop(channel):
    for i in range(messages):
        content = Content(str(i+1))
        send(channel, content)

@defer.inlineCallbacks
def authenticate(client, username, password):
    yield client.authenticate(username, password)
    channel = yield client.channel(1)
    yield channel.channel_open()

    #subscribe
    yield channel.queue_declare(queue='queue', auto_delete=True)
    reply = yield channel.basic_consume(queue='queue')
    queue = yield client.queue(reply.consumer_tag)
    consume_queue(queue)

    # send messages
    reactor.callInThread(send_msgs_from_for_loop, channel)


if __name__ == '__main__':
    import sys
    if len(sys.argv) != 7:
        print "%s host port vhost username password path_to_spec" % sys.argv[0]
        sys.exit(1)

    host = sys.argv[1]
    port = int(sys.argv[2])
    vhost = sys.argv[3]
    username = sys.argv[4]
    password = sys.argv[5]
    specFile = sys.argv[6]

    spec = txamqp.spec.load(specFile)

    delegate = TwistedDelegate()

    d = ClientCreator(reactor, AMQClient, delegate, vhost,
        spec).connectTCP(host, port)
    d.addCallback(authenticate, username, password)
    reactor.run()