#!/usr/bin/env python

from pymongo import MongoClient

from dm.common.utility.loggingManager import LoggingManager
from dm.common.exceptions.dbError import DbError

class DmMongoClient(object):
    """ DM MongoDB client."""

    def __init__(self, dbName, dbUri='mongodb://localhost:27017/'):
        self.logger = LoggingManager.getInstance().getLogger(self.__class__.__name__)
        try:
            self.client = MongoClient(dbUri)
            self.db = self.client[dbName]
            self.dbName = dbName
        except Exception, ex:
            self.logger.error('Cannot connect to Mongo DB: %s' % ex)
            raise DbError(exception=ex)

    def drop(self, collection):
        return self.db[collection].drop()

    def insert(self, collection, dict, **kwargs):
        return self.db[collection].insert(dict, **kwargs)

    def update(self, collection, query, update, **kwargs):
        key = '.'.join((collection, str(update['$set'].keys()[0])))
        return self.db[collection].update(query, update, **kwargs)

    def findOne(self, collection, criteria={}):
        return self.db[collection].find_one(criteria)

    def find(self, collection, criteria={}, projections={}):
        return self.db[collection].find(criteria, projections)

    def findAsList(self, collection, criteria={}, projections={}):
        return list(self.db[collection].find(criteria, projections))

    def getCollectionNames(self):
        return self.db.collection_names()

#######################################################################
# Testing
if __name__ == '__main__':
    mongo = DmMongoClient('dm')
    mongo.drop('students')
    id = mongo.insert('students', { '_id' : 1, 'semester' : 1, 'grades' : [ 70, 87, 90 ]   })
    print 'Student #1 id: ', id
    id = mongo.insert('students', { '_id' : 2, 'semester' : 1, 'grades' : [ 90, 88, 92 ] } )
    print 'Student #2 id: ', id
    print mongo.findAsList('students', criteria={ 'semester' : 1, 'grades': { '$gte' : 85 } }, projections={ 'grades.$' : 1 } )
    result = mongo.update('files', {'_id' : 2}, {'$set' : {'experimentId' : 2}}, upsert=True)
    print result
    doc = mongo.findOne('files', {'experimentId' : 1})
    print doc
    docs = mongo.find('files', projections={'experimentId' : 1})
    print docs.count()
    docs = mongo.findAsList('files', projections={'experimentId' : 1})
    print docs
    print 'Collection names: ', mongo.getCollectionNames()
    mongo.drop('students')
    print 'Collection names (after drop): ', mongo.getCollectionNames()