package com.indeed.proctor.common;

import com.indeed.proctor.common.model.Allocation;
import com.indeed.proctor.common.model.Audit;
import com.indeed.proctor.common.model.ConsumableTestDefinition;
import com.indeed.proctor.common.model.TestBucket;

import javax.annotation.Nonnull;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.TreeMap;

import static com.google.common.collect.Maps.newHashMapWithExpectedSize;
import static java.util.Collections.emptyMap;

/**
 * Return value from {@link Proctor#determineTestGroups(Identifiers, java.util.Map, java.util.Map)}
 * @author ketan
 *
 */
public class ProctorResult {
    public static final ProctorResult EMPTY = new ProctorResult(
            Audit.EMPTY_VERSION,
            emptyMap(),
            emptyMap(),
            emptyMap()
    );

    private final String matrixVersion;
    /**
     * maps from testname to bucket
     */
    @Nonnull
    private final Map<String, TestBucket> buckets;
    /**
     * maps from testname to allocation
     */
    @Nonnull
    private final Map<String, Allocation> allocations;
    /**
     * maps from testname to TestDefinition
     */
    @Nonnull
    private final Map<String, ConsumableTestDefinition> testDefinitions;

    @Deprecated
    public ProctorResult(
            final int matrixVersion,
            @Nonnull final Map<String, TestBucket> buckets,
            @Nonnull final Map<String, ConsumableTestDefinition> testDefinitions
    ) {
        this(Integer.toString(matrixVersion), buckets, emptyMap(), testDefinitions);
    }

    @Deprecated
    public ProctorResult(
            final String matrixVersion,
            @Nonnull final Map<String, TestBucket> buckets,
            @Nonnull final Map<String, ConsumableTestDefinition> testDefinitions
    ) {
        this(matrixVersion, buckets, emptyMap(), testDefinitions);
    }

    public ProctorResult(
            final String matrixVersion,
            @Nonnull final Map<String, TestBucket> buckets,
            @Nonnull final Map<String, Allocation> allocations,
            @Nonnull final Map<String, ConsumableTestDefinition> testDefinitions
    ) {
        this.matrixVersion = matrixVersion;
        this.buckets = new TreeMap<>(buckets);
        this.allocations = new TreeMap<>(allocations);
        this.testDefinitions = new HashMap<>(testDefinitions);
    }

    @SuppressWarnings("UnusedDeclaration")
    public String getMatrixVersion() {
        return matrixVersion;
    }

    @Nonnull
    public Map<String, TestBucket> getBuckets() {
        return buckets;
    }

    @Nonnull
    public Map<String, Allocation> getAllocations() {
        return allocations;
    }

    @Nonnull
    public Map<String, ConsumableTestDefinition> getTestDefinitions() {
        return testDefinitions;
    }

    @Nonnull
    public Map<String, String> getTestVersions() {
        // TODO ImmutableMap?
        final Map<String, String> testVersions = newHashMapWithExpectedSize(testDefinitions.size());
        for (final Entry<String, ConsumableTestDefinition> entry : testDefinitions.entrySet()) {
            testVersions.put(entry.getKey(), entry.getValue().getVersion());
        }
        return testVersions;
    }
}

@@ -276,16 +276,50 @@
 onnull;%0A
+import javax.annotation.Nullable;%0A
 import j
@@ -1439,38 +1439,39 @@
 ,%0A            @N
-onn
 ull
+able
  final Map%3CStrin
@@ -1763,38 +1763,39 @@
 ,%0A            @N
-onn
 ull
+able
  final Map%3CStrin
@@ -2119,32 +2119,85 @@
             
-@Nonnull
+// allowing null for historical reasons%0A            @Nullable
  final Map%3CS
@@ -2395,24 +2395,24 @@
 locations);%0A
-
         this
@@ -2425,24 +2425,65 @@
 efinitions =
+ (testDefinitions == null) ? emptyMap() :
  new HashMap

Make ProctorResult flexible to accomodate some existing unit-tests at Indeed