/*
 *
 *  Copyright 2017 Netflix, Inc.
 *
 *     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 com.netflix.hollow.api.producer;

import com.netflix.hollow.api.consumer.HollowConsumer.BlobRetriever;
import com.netflix.hollow.api.producer.HollowProducer.Populator;
import com.netflix.hollow.api.producer.HollowProducer.ReadState;
import com.netflix.hollow.api.producer.HollowProducer.WriteState;
import com.netflix.hollow.core.index.HollowPrimaryKeyIndex;
import com.netflix.hollow.core.read.engine.HollowReadStateEngine;
import com.netflix.hollow.core.read.engine.HollowTypeReadState;
import com.netflix.hollow.core.schema.HollowObjectSchema;
import com.netflix.hollow.core.schema.HollowSchema;
import com.netflix.hollow.core.schema.HollowSchema.SchemaType;
import com.netflix.hollow.core.util.SimultaneousExecutor;
import com.netflix.hollow.core.write.HollowTypeWriteState;
import com.netflix.hollow.core.write.objectmapper.RecordPrimaryKey;
import com.netflix.hollow.tools.traverse.TransitiveSetTraverser;
import java.util.BitSet;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 
 * Warning: This is a BETA API and is subject to breaking changes.
 * 
 */
public class HollowIncrementalProducer {
    
    private static final Object DELETE_RECORD = new Object();

    private final HollowProducer producer;
    private final ConcurrentHashMap<RecordPrimaryKey, Object> mutations;
    private final double threadsPerCpu;
    
    public HollowIncrementalProducer(HollowProducer producer) {
        this(producer, 1.0d);
    }
    
    public HollowIncrementalProducer(HollowProducer producer, double threadsPerCpu) {
        this.producer = producer;
        this.threadsPerCpu = threadsPerCpu;
        this.mutations = new ConcurrentHashMap<RecordPrimaryKey, Object>();
    }
    
    public void restore(long versionDesired, BlobRetriever blobRetriever) {
        producer.hardRestore(versionDesired, blobRetriever);
    }
    
    public void addOrModify(Object obj) {
        RecordPrimaryKey pk = producer.getObjectMapper().extractPrimaryKey(obj);
        mutations.put(pk, obj);
    }
    
    public void delete(Object obj) {
        RecordPrimaryKey pk = producer.getObjectMapper().extractPrimaryKey(obj);
        delete(pk);
    }
    
    public void delete(RecordPrimaryKey key) {
        mutations.put(key, DELETE_RECORD);
    }

    public static HollowIncrementalProducer.Builder withProducer(HollowProducer producer) {
        HollowIncrementalProducer.Builder builder = new HollowIncrementalProducer.Builder();
        return builder.withProducer(producer);
    }
    
    public static class Builder {
        private HollowProducer producer;
        private double threadsPerCpu = 1.0d; //Default to 1 thread per CPU

        public Builder withProducer(HollowProducer hollowProducer) {
            this.producer = hollowProducer;
            return this;
        }

        public Builder withThreadsPerCpu(double threadsPerCpu) {
            this.threadsPerCpu = threadsPerCpu;
            return this;
        }

        public HollowIncrementalProducer build() {
            if(producer == null)
                throw new IllegalArgumentException("HollowProducer must be specified");

            return new HollowIncrementalProducer(producer, threadsPerCpu);
        }
    }
    
    public long runCycle() {
        return producer.runCycle(new Populator() {
            public void populate(WriteState newState) throws Exception {
                newState.getStateEngine().addAllObjectsFromPreviousCycle();
                removeRecords(newState);
                addRecords(newState);
            }

            private void removeRecords(WriteState newState) {
                Map<String, BitSet> recordsToRemove = findTypesWithRemovedRecords(newState.getPriorState());
                markRecordsToRemove(newState.getPriorState(), recordsToRemove);
                removeRecordsFromNewState(newState, recordsToRemove);
            }

            private Map<String, BitSet> findTypesWithRemovedRecords(ReadState readState) {
                Map<String, BitSet> recordsToRemove = new HashMap<String, BitSet>();
                for(RecordPrimaryKey key : mutations.keySet()) {
                    if(!recordsToRemove.containsKey(key.getType())) {
                        HollowTypeReadState typeState = readState.getStateEngine().getTypeState(key.getType());
                        if(typeState != null) {
                            BitSet bs = new BitSet(typeState.getPopulatedOrdinals().length());
                            recordsToRemove.put(key.getType(), bs);
                        }
                    }
                }
                return recordsToRemove;
            }

            private void markRecordsToRemove(ReadState priorState, Map<String, BitSet> recordsToRemove) {
                HollowReadStateEngine priorStateEngine = priorState.getStateEngine();
                
                for(Map.Entry<String, BitSet> removalEntry : recordsToRemove.entrySet()) {
                    markTypeRecordsToRemove(priorStateEngine, removalEntry.getKey(), removalEntry.getValue());
                }
                
                TransitiveSetTraverser.addTransitiveMatches(priorStateEngine, recordsToRemove);
                TransitiveSetTraverser.removeReferencedOutsideClosure(priorStateEngine, recordsToRemove);
            }

            private void markTypeRecordsToRemove(HollowReadStateEngine priorStateEngine, String type, BitSet typeRecordsToRemove) {
                HollowTypeReadState priorReadState = priorStateEngine.getTypeState(type);
                HollowSchema schema = priorReadState.getSchema();
                if(schema.getSchemaType() == SchemaType.OBJECT) {
                    HollowPrimaryKeyIndex idx = new HollowPrimaryKeyIndex(priorStateEngine, ((HollowObjectSchema) schema).getPrimaryKey()); ///TODO: Should we scan instead?  Can we create this once and do delta updates?
                    
                    for(Map.Entry<RecordPrimaryKey, Object> entry : mutations.entrySet()) {
                        if(entry.getKey().getType().equals(type)) {
                            int priorOrdinal = idx.getMatchingOrdinal(entry.getKey().getKey());
                            
                            if(priorOrdinal != -1)
                                typeRecordsToRemove.set(priorOrdinal);
                        }
                    }
                }
            }
            
            private void removeRecordsFromNewState(WriteState newState, Map<String, BitSet> recordsToRemove) {
                for(Map.Entry<String, BitSet> removalEntry : recordsToRemove.entrySet()) {
                    HollowTypeWriteState writeState = newState.getStateEngine().getTypeState(removalEntry.getKey());
                    BitSet typeRecordsToRemove = removalEntry.getValue();
                    
                    int ordinalToRemove = typeRecordsToRemove.nextSetBit(0);
                    while(ordinalToRemove != -1) {
                        writeState.removeOrdinalFromThisCycle(ordinalToRemove);
                        ordinalToRemove = typeRecordsToRemove.nextSetBit(ordinalToRemove+1);
                    }
                }
            }
            
            private void addRecords(final WriteState newState) {
                SimultaneousExecutor executor = new SimultaneousExecutor(threadsPerCpu);
                for(final Map.Entry<RecordPrimaryKey, Object> entry : mutations.entrySet()) {
                    executor.execute(new Runnable() {
                        public void run() {
                            if(entry.getValue() != DELETE_RECORD)
                                newState.add(entry.getValue());
                        }
                    });
                }

                try {
                    executor.awaitSuccessfulCompletion();
                } catch(Exception e) {
                    throw new RuntimeException(e);
                }
            }

        });
    }

}

@@ -2129,37 +2129,40 @@
     %7D%0A    %0A    p
-ublic
+rotected
  HollowIncrement

Make the HollowIncrementalProducer constructor protected