/*
 * Copyright (C) 2011 Benoit GUEROUT <bguerout at gmail dot com> and Yves AMSELLEM <amsellem dot yves at gmail dot com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.jongo;

import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.WriteConcern;
import com.mongodb.WriteResult;
import org.bson.types.ObjectId;
import org.jongo.bson.BsonDocument;
import org.jongo.marshall.Marshaller;

class Save {

    private final Marshaller marshaller;
    private final DBCollection collection;
    private final ObjectIdUpdater objectIdUpdater;
    private final Object pojo;
    private WriteConcern concern;

    Save(DBCollection collection, Marshaller marshaller, ObjectIdUpdater objectIdUpdater, Object pojo) {
        this.marshaller = marshaller;
        this.collection = collection;
        this.objectIdUpdater = objectIdUpdater;
        this.pojo = pojo;
    }

    public Save concern(WriteConcern concern) {
        this.concern = concern;
        return this;
    }

    public WriteResult execute() {
        DBObject dbObject;
        if (objectIdUpdater.canSetObjectId(pojo)) {
            ObjectId id = ObjectId.get();
            objectIdUpdater.setDocumentGeneratedId(pojo, id);
            dbObject = marshallDocument().toDBObject();
            dbObject.put("_id", id);
        } else {
            dbObject = marshallDocument().toDBObject();
        }

        return collection.save(dbObject, determineWriteConcern());
    }

    private BsonDocument marshallDocument() {
        try {
            return marshaller.marshall(pojo);
        } catch (Exception e) {
            String message = String.format("Unable to save object %s due to a marshalling error", pojo);
            throw new IllegalArgumentException(message, e);
        }
    }

    private WriteConcern determineWriteConcern() {
        return concern == null ? collection.getWriteConcern() : concern;
    }
}

@@ -722,114 +722,43 @@
 odb.
-DBCollection;%0Aimport com.mongodb.DBObject;%0Aimport com.mongodb.WriteConcern;%0Aimport com.mongodb.WriteResult
+*;%0Aimport org.bson.LazyBSONCallback
 ;%0Aim
@@ -862,16 +862,65 @@
 aller;%0A%0A
+import java.util.HashSet;%0Aimport java.util.Set;%0A%0A
 class Sa
@@ -1626,342 +1626,683 @@
     
+db
 Object
-Id id = ObjectId.get();%0A            objectIdUpdater.setDocumentGeneratedId(pojo, id);%0A            dbObject = marshallDocument().toDBObject(
+ = createDBObjectToInsert();%0A        %7D else %7B%0A            dbObject = createDBObjectToUpdate();%0A        %7D%0A%0A        return collection.save(dbObject, determineWriteConcern());%0A    %7D%0A%0A    private DBObject createDBObjectToUpdate() %7B%0A        BsonDocument document = marshallDocument();%0A        return new UncheckableDBObject(document.toByteArray()
 );%0A    
+%7D%0A%0A
     
-    dbObject.put(%22_id%22, id);%0A        %7D else %7B%0A            dbObject = marshallDocument().toDBObject();%0A        %7D%0A%0A        return collection.save(dbObject, determineWriteConcern())
+private DBObject createDBObjectToInsert() %7B%0A        ObjectId id = ObjectId.get();%0A        objectIdUpdater.setDocumentGeneratedId(pojo, id);%0A%0A        BsonDocument document = marshallDocument();%0A        DBObject dbo = new UncheckableDBObject(document.toByteArray());%0A        dbo.put(%22_id%22, id);%0A%0A        return dbo
 ;%0A  
@@ -2755,11 +2755,390 @@
 n;%0A    %7D
+%0A%0A    private static class UncheckableDBObject extends LazyWriteableDBObject %7B%0A%0A        private final HashSet%3CString%3E keys;%0A%0A        private UncheckableDBObject(byte%5B%5D data) %7B%0A            super(data, new LazyBSONCallback());%0A            this.keys = new HashSet%3CString%3E();%0A        %7D%0A%0A        @Override%0A        public Set%3CString%3E keySet() %7B%0A            return keys;%0A        %7D%0A    %7D
 %0A%7D%0A

Improve save speed by passing a fake DBObject