#!/usr/bin/env python
"""
Subscribe to an AMQP exchange and dump the messages.
"""
# standard
import sys
import optparse
import time
# third party
# local
from netlogger.parsers.base import NLSimpleParser
from netlogger.amqp.amqp_process import TableBuilder, rpython, bson_decode

class Performance:
    def __init__(self):
        self.count, self.i = 0, 0

    def start(self):
        self.t0 = time.time()
        
    def item(self):
        self.count += 1
        self.i += 1
        if self.i == 1000:
            t1 = time.time()
            print("{count:d}: {rate:g} events/sec\n".format(
                count=self.count, rate=self.count / (t1 - self.t0)))
            self.i = 0

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]
    op = optparse.OptionParser(usage="%prog [options]", description=__doc__)
    op.add_option("--host", dest="host", default="localhost",
                  help="AMQP broker host (default=%default)")
    op.add_option("--port", dest="port", type="int", default=5672,
                  help="AMQP broker port (default=%default)")
    op.add_option("--exchange", dest="exch", default="amq.{queue_type}",
                  help="Exchange name (default=%default)")
    op.add_option("--type", dest="type", default="direct",
                  help="Exchange type: direct, topic, or fanout "
                  "(default=%default)")
    op.add_option("--queue", dest="queue", default="testq",
                  help="Queue name (default=%default)")
    op.add_option("--route", dest="route", metavar="key", default="#",
                  help="Routing key (default=%default)")    
    op.add_option("--durable", dest="dur", action="store_true", default=False,
                  help="Is exchange durable? (default=no)")
    op.add_option("--bson", dest="is_bson", action="store_true", default=False,
                  help="Are messages BSON? (default=no; BP)")
    op.add_option("--null", dest="is_null", action="store_true", default=False,
                  help="Suppress output to stdout (default=no)")
    op.add_option("--perf", dest="perf", action="store_true", default=False,
                  help="Print performance information to stdout (default=no)")
    opts, args = op.parse_args(argv)
    # set parsing
    if opts.is_bson:
        if bson_decode is None:
            op.error("BSON specified but unable to load bson module")
        parse = bson_decode
    else:
        parse = NLSimpleParser().parseLine
    # connect to amqp
    autodel = True
    if opts.exch == "amq.{queue_type}":
        opts.exch = opts.exch.format(queue_type=opts.type)
        autodel = False
    kw = dict(host=opts.host, port=opts.port, queue=opts.queue,
              exchange=opts.exch, exchange_type=opts.type, durable=opts.dur,
              routing_key=opts.route, auto_delete=autodel)
    if opts.perf:
        perf = Performance()
    else:
        perf = None
    output = 'r' # XXX
    do_print = not opts.is_null
    if output == 'csv':
        tb = TableBuilder(to_dict=parse, required_attr=["ts", "event"],
                          optional_attr=["color", "size"], last={"event":"stage.end"}, **kw)
        if perf:
            perf.start()
        for (body, was_proc) in tb:
            #print("got message, processed={wp}".format(wp=str(was_proc)))
            if do_print:
                tb.write_csv(csv.writer(sys.stdout), hdr=True)
            if perf:
                perf.item()
    elif output == 'r':
        if rpython is None:
            op.error("R mode selected but cannot import R wrapper library")
        tb = TableBuilder(to_dict=parse, required_attr=["ts", "event"],
                          optional_attr=["color", "size"], last={"event":"stage.end"}, **kw)
        if perf:
            perf.start()
        nrows = 0
        for (body, was_proc) in tb:
            if was_proc:
                nrows += 1
                if nrows == 100:
                    df = tb.get_rdata()
                    nrows = 0
                perf.item()
                
    elif output == 'raw':
        reader = Reader(**kw)
        # main loop
        if perf:
            perf.start()
        for item in reader:
            if do_print:
                print(parse(item))
            if perf:
                perf.item()
    # done
    return 0

if __name__ == '__main__':
    sys.exit(main())
