#!/usr/bin/env python

from bson.objectid import ObjectId
from dm.common.utility.loggingManager import LoggingManager
from dm.common.exceptions.invalidArgument import InvalidArgument
from dm.common.exceptions.objectAlreadyExists import ObjectAlreadyExists
from dm.common.exceptions.objectNotFound import ObjectNotFound
from dm.common.exceptions.dbError import DbError

class DmMongoCollection(object):
    """Base collection in mongo db."""

    ALL_FIELDS_DICT = {'__return_only_id__' : False}
    UNIQUE_KEYS_LIST = []

    def __init__(self, collectionName, dbClient):
        self.logger = LoggingManager.getInstance().getLogger(self.__class__.__name__)
        self.dbClient = dbClient

        # Collection consists of items: [item]
        self.collectionName = collectionName

        # items => item
        self.itemName = collectionName[:-1] 

        # item => Item
        self.capitalizedItemName = self.itemName.capitalize() 

    @classmethod
    def getUniqueKeys(cls):
        return cls.UNIQUE_KEYS_LIST

    def findByKey(self, key, value):
        dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
        if dbObjectDict is None:
            raise ObjectNotFound('%s with %s=%s not found.' % (self.capitalizedItemName, key, value))
        return dbObjectDict

    def findByKeys(self, keyList, queryDict):
        for key in keyList:
            if not queryDict.has_key(key):
                raise InvalidArgument('%s query dictionary does not specify key %s.' % (self.capitalizedItemName, key))

        dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
        if dbObjectDict is None:
            raise ObjectNotFound('%s with properties %s not found.' % (self.capitalizedItemName, queryDict))
        return dbObjectDict

    def findByUniqueKeys(self, queryDict):
        return self.findByKeys(self.UNIQUE_KEYS_LIST, queryDict)

    def findById(self, id):
        return self.findByKey('_id', ObjectId(id))

    def findByName(self, name):
        return self.findByKey(self.NAME_KEY, name)

    def findByQueryDict(self, queryDict, returnFieldDict=ALL_FIELDS_DICT):
        return self.dbClient.findAsList(self.collectionName, queryDict, returnFieldDict)

    def listByKey(self, key):
        return self.dbClient.findAsList(self.collectionName, {}, {key : True})

    def listById(self):
        return self.dbClient.findAsList(self.collectionName, {}, {})

    def listByName(self):
        return self.findByName('_name')

    def listByKeys(self, keyList):
        returnFieldDict = {}
        for key in keyList:
            returnFieldDict[key] = True
        return self.dbClient.findAsList(self.collectionName, {}, returnFieldDict)
    def listAll(self):
        return self.dbClient.findAsList(self.collectionName, {}, DmMongoCollection.ALL_FIELDS_DICT)

    def __addDbObject(self, objectDict):
        try:
            self.dbClient.insert(self.collectionName, objectDict)
        except Exception, ex:
            self.logger.error('Cannot add %s with %s set to %s: %s' % (self.itemName, key, value, ex))
            raise DbError(exception=ex)
        return objectDict

    def addByKey(self, key, objectDict):
        value = objectDict.get(key)
        if value is None:
            raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))

        dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
        if dbObjectDict is not None:
            raise ObjectAlreadyExists('%s with %s set to %s already exists.' % (self.capitalizedItemName, key, value))

        return self.__addDbObject(objectDict)

    def addByName(self, objectDict):
        return self.addByKey('_name', objectDict)

    def addByKeys(self, keyList, objectDict):
        queryDict = {}
        for key in keyList:
            value = objectDict.get(key)
            if value is None:
                raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
            queryDict[key] = value

        dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
        if dbObjectDict is not None:
            raise ObjectAlreadyExists('%s with properties %s already exists.' % (self.capitalizedItemName, queryDict))

        return self.__addDbObject(objectDict)

    def addByUniqueKeys(self, objectDict):
        return self.addByKeys(self.UNIQUE_KEYS_LIST, objectDict)

    def __updateDbObject(self, dbObjectDict, objectDict):
        try:
            id = dbObjectDict.get('_id')
            dbObjectDict.update(objectDict)
            self.dbClient.update(self.collectionName,
                {'_id' : id},
                {'$set' : dbObjectDict})
        except Exception, ex:
            self.logger.error('Cannot update %s %s with %s: %s' % (self.itemName, dbObjectDict, objectDict, ex))
            raise DbError(exception=ex)
        return dbObjectDict

    def updateByKey(self, key, objectDict):
        value = objectDict.get(key)
        if value is None:
            raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))

        dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
        if dbObjectDict is None:
            raise ObjectNotFound('%s with %s set to %s not found.' % (self.capitalizedItemName, key, value))

        return self.__updateDbObject(dbObjectDict, objectDict)

    def updateByKeys(self, keyList, objectDict):
        queryDict = {}
        for key in keyList:
            value = objectDict.get(key)
            if value is None:
                raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
            queryDict[key] = value

        dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
        if dbObjectDict is None:
            raise ObjectNotFound('%s with properties %s not found.' % (self.capitalizedItemName, queryDict))

        return self.__updateDbObject(dbObjectDict, objectDict)

    def updateByUniqueKeys(self, objectDict):
        return self.updateByKeys(self.UNIQUE_KEYS_LIST, objectDict)

    def updateById(self, objectDict):
        # Convert 'id' to '_id' key if needed, and wrap it with ObjectId()
        if not objectDict.has_key('_id') and objectDict.has_key('id'):
            objectDict['_id'] = ObjectId(objectDict['id'])
            del objectDict['id']
        return self.updateByKey('_id', objectDict)

    def updateByName(self, objectDict):
        return self.updateByKey('_name', objectDict)

    def updateOrAddByKey(self, key, objectDict):
        value = objectDict.get(key)
        if value is None:
            raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))

        dbObjectDict = self.dbClient.findOne(self.collectionName, {key : value})
        if dbObjectDict is None:
            return self.__addDbObject(objectDict)
        else:
            return self.__updateDbObject(dbObjectDict, objectDict)

    def updateOrAddByKeys(self, keyList, objectDict):
        queryDict = {}
        for key in keyList:
            value = objectDict.get(key)
            if value is None:
                raise InvalidArgument('%s info dictionary does not specify key %s.' % (self.capitalizedItemName, key))
            queryDict[key] = value

        dbObjectDict = self.dbClient.findOne(self.collectionName, queryDict)
        if dbObjectDict is None:
            return self.__addDbObject(objectDict)
        else:
            return self.__updateDbObject(dbObjectDict, objectDict)

    def updateOrAddByUniqueKeys(self, objectDict):
        return self.updateOrAddByKeys(self.UNIQUE_KEYS_LIST, objectDict)

    def updateOrAddByName(self, objectDict):
        return self.updateOrAddByKey('_name', objectDict)

#######################################################################
# Testing
if __name__ == '__main__':
    from dmMongoClient import DmMongoClient
    from bson.objectid import ObjectId
    mongo = DmMongoClient('dm')
    fileCollectionImpl = DmMongoCollection('files', mongo)
    #objectDict = {'name' : 'xyz-001', 'experiment' : 'myexp-D', 'update' : 'sv2', 'locationList' : '[/opt/xyz, /data/xyz]'}
    #print fileCollectionImpl.updateOrAddByKey('name', objectDict)
    #f = fileCollectionImpl.findByKey('name', 'xyz-001')
    #print f
    #print type(f['_id'])
    #print fileCollectionImpl.findById('556de0059e058b0ef4c4413b')
    #print fileCollectionImpl.findByQueryDict({'experiment' : 'exp-001'}, {'locationList' : 1})

    #print 'LIST BY ID\n', fileCollectionImpl.listById()
    #print 'LIST BY NAME\n', fileCollectionImpl.listByKey('name')
    print 'LIST ALL\n'
    for f in fileCollectionImpl.listAll():
        print 'FILE: %s\n' % f


    #print fileCollectionImpl.updateByKeys(['name', 'experiment'], objectDict)