from bson.objectid import ObjectId

from girder.models.model_base import AccessControlledModel
from girder.constants import AccessType

from .molecule import Molecule as MoleculeModel


class Geometry(AccessControlledModel):

    def __init__(self):
        super(Geometry, self).__init__()

    def initialize(self):
        self.name = 'geometry'
        self.ensureIndex('moleculeId')

        self.exposeFields(level=AccessType.READ, fields=(
            '_id', 'moleculeId', 'cjson', 'provenanceType', 'provenanceId'))

    def validate(self, doc):
        # If we have a moleculeId ensure it is valid.
        if 'moleculeId' in doc:
            mol = MoleculeModel().load(doc['moleculeId'], force=True)
            doc['moleculeId'] = mol['_id']

        return doc

    def create(self, user, moleculeId, cjson, provenanceType=None,
               provenanceId=None, public=False):
        geometry = {
            'moleculeId': moleculeId,
            'cjson': cjson,
            'creatorId': user['_id']
        }

        if provenanceType is not None:
            geometry['provenanceType'] = provenanceType

        if provenanceId is not None:
            geometry['provenanceId'] = provenanceId

        self.setUserAccess(geometry, user=user, level=AccessType.ADMIN)
        if public:
            self.setPublic(geometry, True)

        return self.save(geometry)

    def find_geometries(self, moleculeId, user=None):
        query = {
            'moleculeId': ObjectId(moleculeId)
        }

        return self.findWithPermissions(query, user=user)

@@ -134,16 +134,32 @@
 e%0A%0Afrom 
+molecules.models
 .molecul
@@ -193,16 +193,79 @@
 leModel%0A
+from molecules.utilities.whitelist_cjson import whitelist_cjson
 %0A%0Aclass 
@@ -961,16 +961,90 @@
 =False):
+%0A%0A        # We will whitelist the cjson to only include the geometry parts
 %0A       
@@ -1116,21 +1116,38 @@
 cjson': 
-cjson
+whitelist_cjson(cjson)
 ,%0A      

Whitelist the cjson when geometries are created