import java.io.PrintStream;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import org.checkerframework.checker.nullness.qual.EnsuresNonNullIf;
import org.checkerframework.checker.nullness.qual.KeyFor;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

public class NonNullMapValue {

    // Discussion:
    //
    // It can be useful to indicate that all the values in a map are non-null.
    // ("@NonNull" is redundant in this declaration, but I've written it
    // explicitly because it is the annotation we are talking about.)
    //   HashMap<String,@NonNull String> myMap;
    //
    // However, the get method's declaration is misleading (in the context of
    // the nullness type system), since it can always return null no matter
    // whether the map values are non-null:
    //   V get(Object key) { ... return null; }
    //
    // Here are potential solutions:
    //  * Forbid declaring values as non-null.  This is the wrong approach.
    //    (It would also be hard to express syntactically.)
    //  * The checker could recognize a special new annotation on the return
    //    value of get, indicating that its return type isn't merely inferred
    //    from the generic type, but is always nullable.  (This special new
    //    annotations could even be "@Nullable".  A different annotation may
    //    be better, becase in general we would like to issue an error
    //    message when someone applies an annotation to a generic type
    //    parameter.)
    // Additionally, to reduce the number of false positive warnings caused
    // by the fact that get's return value is nullable:
    //  * Build a more specialized sophisticated flow analysis that checks
    //    that the passed key to Map.containsKey() is either checked against
    //    Map.containsKey() or Map.keySet().

    Map<String, @NonNull String> myMap;

    NonNullMapValue(Map<String, @NonNull String> myMap) {
        this.myMap = myMap;
    }

    void testMyMap(String key) {
        @NonNull String value;
        // :: error: (assignment.type.incompatible)
        value = myMap.get(key); // should issue warning
        if (myMap.containsKey(key)) {
            value = myMap.get(key);
        }
        for (String keyInMap : myMap.keySet()) {
            // :: error: (assignment.type.incompatible)
            value = myMap.get(key); // should issue warning
        }
        for (String keyInMap : myMap.keySet()) {
            value = myMap.get(keyInMap);
        }
        for (Map.Entry<@KeyFor("myMap") String, @NonNull String> entry : myMap.entrySet()) {
            String keyInMap = entry.getKey();
            value = entry.getValue();
        }
        for (Iterator<@KeyFor("myMap") String> iter = myMap.keySet().iterator(); iter.hasNext(); ) {
            String keyInMap = iter.next();
            // value = myMap.get(keyInMap);
        }
        value = myMap.containsKey(key) ? myMap.get(key) : "hello";
    }

    public static <T> void print(Map<T, List<T>> graph, PrintStream ps, int indent) {
        for (T node : graph.keySet()) {
            for (T child : graph.get(node)) {
                ps.printf("  %s%n", child);
            }
            @NonNull List<T> children = graph.get(node);
            for (T child : children) {
                ps.printf("  %s%n", child);
            }
        }
    }

    public static <T> void testAssertFlow(Map<T, List<T>> preds, T node) {
        assert preds.containsKey(node);
        for (T pred : preds.get(node)) {}
    }

    public static <T> void testContainsKey1(Map<T, List<T>> dom, T pred) {
        assert dom.containsKey(pred);
        // Both of the next two lines should type-check.  The second one won't
        // unless the checker knows that pred is a key in the map.
        List<T> dom_of_pred1 = dom.get(pred);
        @NonNull List<T> dom_of_pred2 = dom.get(pred);
    }

    public static <T> void testContainsKey2(Map<T, List<T>> dom, T pred) {
        if (!dom.containsKey(pred)) {
            throw new Error();
        }
        // Both of the next two lines should type-check.  The second one won't
        // unless the checker knows that pred is a key in the map.
        List<T> dom_of_pred1 = dom.get(pred);
        @NonNull List<T> dom_of_pred2 = dom.get(pred);
    }

    public static void process_unmatched_procedure_entries() {
        HashMap<Integer, Date> call_hashmap = new HashMap<>();
        for (Integer i : call_hashmap.keySet()) {
            @NonNull Date d = call_hashmap.get(i);
        }
        Set<@KeyFor("call_hashmap") Integer> keys = call_hashmap.keySet();
        for (Integer i : keys) {
            @NonNull Date d = call_hashmap.get(i);
        }
        Set<@KeyFor("call_hashmap") Integer> keys_sorted =
                new TreeSet<@KeyFor("call_hashmap") Integer>(call_hashmap.keySet());
        for (Integer i : keys_sorted) {
            @NonNull Date d = call_hashmap.get(i);
        }
    }

    public static Object testPut(Map<Object, Object> map, Object key) {
        if (!map.containsKey(key)) {
            map.put(key, new Object());
        }
        return map.get(key);
    }

    public static Object testAssertGet(Map<Object, Object> map, Object key) {
        assert map.get(key) != null;
        return map.get(key);
    }

    public static Object testThrow(Map<Object, Object> map, Object key) {
        if (!map.containsKey(key)) {
            if (true) {
                return "m";
            } else {
                throw new RuntimeException();
            }
        }
        return map.get(key);
    }

    public void negateMap(Map<Object, Object> map, Object key) {
        if (!map.containsKey(key)) {
        } else {
            @NonNull Object v = map.get(key);
        }
    }

    public void withinElseInvalid(Map<Object, Object> map, Object key) {
        if (map.containsKey(key)) {
        } else {
            // :: error: (assignment.type.incompatible)
            @NonNull Object v = map.get(key); // should issue warning
        }
    }

    // Map.get should be annotated as @org.checkerframework.dataflow.qual.Pure
    public static int mapGetSize(MyMap<Object, List<Object>> covered, Object file) {
        return (covered.get(file) == null) ? 0 : covered.get(file).size();
    }

    interface MyMap<K, V> extends Map<K, V> {
        // TODO: @AssertGenericNullnessIfTrue("get(#1)")
        @org.checkerframework.dataflow.qual.Pure
        public abstract boolean containsKey(@Nullable Object a1);

        // We get an override warning, because we do not use the annotated JDK in the
        // test suite. Ignore this.
        @SuppressWarnings("override.return.invalid")
        @org.checkerframework.dataflow.qual.Pure
        public @Nullable V get(@Nullable Object o);
    }

    private static final String KEY = "key";
    private static final String KEY2 = "key2";

    void testAnd(MyMap<String, String> map, MyMap<String, @Nullable String> map2) {
        if (map.containsKey(KEY)) {
            map.get(KEY).toString();
        }
        // :: warning: (known.nonnull)
        if (map.containsKey(KEY2) && map.get(KEY2).toString() != null) {}
        // :: error: (dereference.of.nullable) :: warning: (known.nonnull)
        if (map2.containsKey(KEY2) && map2.get(KEY2).toString() != null) {}
    }

    void testAndWithIllegalMapAnnotation(MyMap2<String, String> map) {
        if (map.containsKey(KEY)) {
            map.get(KEY).toString();
        }
        // :: warning: (known.nonnull)
        if (map.containsKey(KEY2) && map.get(KEY2).toString() != null) {
            // do nothing
        }
    }

    interface MyMap2<K, V> {
        @org.checkerframework.dataflow.qual.Pure
        // This annotation is not legal on containsKey in general.
        // If the Map is declared as (say) Map<Object, @Nullable Object>,
        // then get returns a nullable value.  We really want to say that if
        // containsKey returns non-null, then get returns V rather than
        // @Nullable V, but I don't know how to say that.
        @EnsuresNonNullIf(result = true, expression = "get(#1)")
        public abstract boolean containsKey(@Nullable Object a1);

        @org.checkerframework.dataflow.qual.Pure
        public abstract @Nullable V get(@Nullable Object a1);
    }

    interface MyMap3<K, V> {
        @org.checkerframework.dataflow.qual.Pure
        @EnsuresNonNullIf(result = true, expression = "get(#1)")
        // The following error is issued because, unlike in interface MyMap2,
        // this interface has no get() method.
        // :: error: (flowexpr.parse.error)
        boolean containsKey(@Nullable Object a1);
    }
}