/**
 * Copyright 2015 Palantir Technologies
 *
 * Licensed under the BSD-3 License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://opensource.org/licenses/BSD-3-Clause
 *
 * 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.palantir.atlasdb.keyvalue.partition;

import java.util.Map;
import java.util.concurrent.Future;

import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import com.palantir.atlasdb.keyvalue.partition.QuorumParameters.QuorumRequestParameters;

public class QuorumTracker<FutureReturnType, TrackingUnit> {

    private final Map<TrackingUnit, Integer> numberOfRemainingSuccessesForSuccess;
    private final Map<TrackingUnit, Integer> numberOfRemainingFailuresForFailure;
    private final Map<Future<FutureReturnType>, Iterable<TrackingUnit>> unitsByReference;
    private boolean failure;

    public QuorumTracker(Iterable<TrackingUnit> allTrackedUnits,
                         QuorumRequestParameters quorumRequestParameters) {

        numberOfRemainingFailuresForFailure = Maps.newHashMap();
        numberOfRemainingSuccessesForSuccess = Maps.newHashMap();

        for (TrackingUnit item : allTrackedUnits) {
            numberOfRemainingSuccessesForSuccess.put(item, quorumRequestParameters.getSuccessFactor());
            numberOfRemainingFailuresForFailure.put(item, quorumRequestParameters.getFailureFactor());
        }

        unitsByReference = Maps.newHashMap();
        failure = false;
    }

    public static <FutureReturnType, TrackingUnit> QuorumTracker<FutureReturnType, TrackingUnit>
            of(Iterable<TrackingUnit> allTrackedUnits,
                    QuorumRequestParameters quorumRequestParameters) {
        return new QuorumTracker<FutureReturnType, TrackingUnit>(
                allTrackedUnits,
                quorumRequestParameters);
    }

    public void handleSuccess(Future<FutureReturnType> ref) {
        Preconditions.checkState(!finished());
        Preconditions.checkArgument(unitsByReference.containsKey(ref));

        for (TrackingUnit unit : unitsByReference.get(ref)) {
            if (numberOfRemainingSuccessesForSuccess.containsKey(unit)) {
                int newValue = numberOfRemainingSuccessesForSuccess.get(unit) - 1;
                if (newValue == 0) {
                    numberOfRemainingSuccessesForSuccess.remove(unit);
                    numberOfRemainingFailuresForFailure.remove(unit);
                } else {
                    numberOfRemainingSuccessesForSuccess.put(unit, newValue);
                }
            }
        }
        unregisterRef(ref);
    }

    public void handleFailure(Future<FutureReturnType> ref) {
        Preconditions.checkState(!finished());
        Preconditions.checkArgument(unitsByReference.containsKey(ref));
        for (TrackingUnit cell : unitsByReference.get(ref)) {
            if (numberOfRemainingFailuresForFailure.containsKey(cell)) {
                int newValue = numberOfRemainingFailuresForFailure.get(cell) - 1;
                if (newValue == 0) {
                    failure = true;
                    break;
                } else {
                    numberOfRemainingFailuresForFailure.put(cell, newValue);
                }
            }
        }
        unregisterRef(ref);
    }

    public void registerRef(Future<FutureReturnType> ref, Iterable<TrackingUnit> items) {
        Preconditions.checkState(!finished());
        unitsByReference.put(ref, items);
    }

    private void unregisterRef(Future<FutureReturnType> ref) {
        Preconditions.checkArgument(unitsByReference.containsKey(ref));
        unitsByReference.remove(ref);
    }

    public void cancel(boolean mayInterruptIfRunning) {
        for (Future<FutureReturnType> f : unitsByReference.keySet()) {
            f.cancel(mayInterruptIfRunning);
        }
    }

    public boolean failed() {
        return failure;
    }

    public boolean succeeded() {
        return !failed() && numberOfRemainingSuccessesForSuccess.isEmpty();
    }

    public boolean finished() {
        return failed() || succeeded();
    }
}
@@ -863,16 +863,58 @@
 ameters;
+%0Aimport com.palantir.common.collect.Maps2;
 %0A%0Apublic
@@ -1395,17 +1395,16 @@
 ters) %7B%0A
-%0A
         
@@ -1449,199 +1449,48 @@
 Maps
-.newHashMap();%0A        numberOfRemainingSuccessesForSuccess = Maps.newHashMap();%0A%0A        for (TrackingUnit item : allTrackedUnits) %7B%0A            numberOfRemainingSuccessesForSuccess.put(item
+2.createConstantValueMap(allTrackedUnits
 , qu
@@ -1514,23 +1514,23 @@
 ters.get
-Success
+Failure
 Factor()
@@ -1528,28 +1528,24 @@
 eFactor());%0A
-    
         numb
@@ -1561,35 +1561,74 @@
 ning
-FailuresForFailure.put(item
+SuccessesForSuccess = Maps2.createConstantValueMap(allTrackedUnits
 , qu
@@ -1644,39 +1644,39 @@
 stParameters.get
-Failure
+Success
 Factor());%0A     
@@ -1674,19 +1674,8 @@
 ));%0A
-        %7D%0A%0A
     

Use helper methods for constant value maps.