#!/usr/bin/env python

# Script to record a ros topic to the db
#
# Requires db to be running

import roslib
roslib.load_manifest('mongo_ros')
import argparse
import sys
import mongo_ros as mr
import re
import rospy
from mongo_ros import util
from math import log

class TopicSubscriber(object):

    def __init__(self, topic, host, port, db, coll_name, pkg, msg):
        self.num_received = 0
        self.cls = util.load_message_class(pkg, msg)
        self.coll = mr.MessageCollection(db, coll_name, self.cls,
                                         db_host=host, db_port=port)
        self.sub = rospy.Subscriber(topic, rospy.AnyMsg, self.cb)

    def cb(self, raw):
        self.num_received += 1
        l = log(self.num_received)/log(2)
        if abs(self.num_received-pow(2, int(l)))<0.1:
            print "Have received {0} messages".format(self.num_received)
        msg = self.cls()
        msg.deserialize(raw._buff)
        self.coll.insert(msg)
        
def parse_msg(str):
    l = str.split('/')
    if len(l) != 2:
        raise ValueError("Invalid message type {0}".format(str))
    return (l[0], l[1])

def main(argv=sys.argv):
    argv = rospy.myargv(argv)
    parser = argparse.ArgumentParser(description="Record ros topic to db")
    parser.add_argument('topic', help="Ros topic to record")
    parser.add_argument('database', help="Name of db to record to")
    parser.add_argument('msg_type', help="Specified as single string PKG/MSG")
    parser.add_argument('--host', '-o', default='localhost', help="Host for db server")
    parser.add_argument('--port', '-p', default=27017, help="Port for db server", type=int)
    parser.add_argument('--collection', '-c', default='',
                        help="Name of collection within db."
                        "  Defaults to using (last component of) topic name.")
    args = parser.parse_args(argv[1:])

    print "Topic: {0}, db: {1}".format(args.topic, args.database)
    if args.collection:
        coll = args.collection
    else:
        coll = args.topic.lstrip('/').replace("/", ".")
    pkg, msg = parse_msg(args.msg_type)

    rospy.init_node('record_to_db', anonymous=True)
    
    print("Recording {0.topic} of type {2}/{3} to {0.host}:{0.port}:{0.database}/{1}".\
          format(args, coll, pkg, msg))

    sub = TopicSubscriber(args.topic, args.host, args.port, args.database, coll, pkg, msg)
    rospy.spin()
    
if __name__=='__main__':
    sys.exit(main())


        
        

        
                                    
