import java.util.HashMap;
import java.util.Vector;
import org.checkerframework.checker.nullness.qual.*;

public class KeyForFlow extends HashMap<String, Object> {

  String k = "key";
  HashMap<String, Object> m = new HashMap<>();

  void testContainsKeyForLocalKeyAndLocalMap() {
    String k_local = "key";
    HashMap<String, Object> m_local = new HashMap<>();

    if (m_local.containsKey(k_local)) {
      @KeyFor("m_local") Object s = k_local;
    }

    // :: error: (assignment)
    @KeyFor("m_local") String s2 = k_local;
  }

  void testContainsKeyForLocalKeyAndFieldMap() {
    String k_local = "key";

    if (m.containsKey(k_local)) {
      @KeyFor("m") Object s = k_local;
    }

    // :: error: (assignment)
    @KeyFor("m") String s2 = k_local;
  }

  void testContainsKeyForFieldKeyAndLocalMap() {
    HashMap<String, Object> m_local = new HashMap<>();

    if (m_local.containsKey(k)) {
      @KeyFor("m_local") Object s = k;
    }

    // :: error: (assignment)
    @KeyFor("m_local") String s2 = k;
  }

  void testContainsKeyForFieldKeyAndFieldMap() {
    if (m.containsKey(k)) {
      @KeyFor("m") Object s = k;
    }

    // :: error: (assignment)
    @KeyFor("m") String s2 = k;
  }

  static String k_s = "key";

  void testContainsKeyForStaticKeyAndFieldMap() {
    if (m.containsKey(k_s)) {
      @KeyFor("m") Object s = k_s;
    }

    // :: error: (assignment)
    @KeyFor("m") String s2 = k_s;
  }

  static HashMap<String, Object> m_s = new HashMap<>();

  void testContainsKeyForFieldKeyAndStaticMap() {
    if (m_s.containsKey(k)) {
      // Currently for this to work, the user must write @KeyFor("classname.static_field")
      @KeyFor("m_s") Object s = k;
    }

    // :: error: (assignment)
    @KeyFor("m_s") String s2 = k;
  }

  void testContainsKeyForFieldKeyAndReceiverMap() {
    if (containsKey(k)) {
      @KeyFor("this") Object s = k;
    }

    // :: error: (assignment)
    @KeyFor("this") String s2 = k;
  }

  // TODO: The diamond operator does not work here:
  //    Vector<@KeyFor("m2") String> coll = new Vector<>();
  // Figure out why not.
  Vector<@KeyFor("m2") String> coll = new Vector<@KeyFor("m2") String>();
  HashMap<String, Object> m2 = new HashMap<>();
  String k2 = "key2";

  void testCallingPutAfterAdd() {
    // :: error: (argument)
    coll.add(k2);
    m2.put(k2, new Object());
  }

  void testPutForLocalKeyAndLocalMap() {
    HashMap<String, Object> m2_local = new HashMap<>();
    Vector<@KeyFor("m2_local") String> coll_local = new Vector<>();
    String k2_local = "key2";

    m2_local.put(k2_local, new Object());
    coll_local.add(k2_local);
  }

  void testPutForLocalKeyAndFieldMap() {
    String k2_local = "key2";

    m2.put(k2_local, new Object());
    coll.add(k2_local);
  }

  void testPutForFieldKeyAndLocalMap() {
    HashMap<String, Object> m2_local = new HashMap<>();
    Vector<@KeyFor("m2_local") String> coll_local = new Vector<>();

    m2_local.put(k2, new Object());
    coll_local.add(k2);
  }

  void testPutForFieldKeyAndFieldMap() {
    m2.put(k2, new Object());
    coll.add(k2);
  }

  /*
  This scenario is not working since in Vector, "this" gets translated to "coll_local".
  The same thing happens if the collection is a field instead of a local.
  However this seems like a low-priority scenario to enable.

  void testPutForFieldKeyAndReceiverMap() {
    Vector<@KeyFor("this") String> coll_local = new Vector<>();

    put(k2, new Object());
    coll_local.add(k2);
  }*/

  class foo {
    public HashMap<String, Object> m = new HashMap<>();
  }

  void testContainsKeyForFieldKeyAndMapFieldOfOtherClass() {
    foo f = new foo();

    if (f.m.containsKey(k)) {
      @KeyFor("f.m") Object s = k;
    }

    // :: error: (assignment)
    @KeyFor("f.m") String s2 = k;
  }

  void testPutForFieldKeyAndMapFieldOfOtherClass() {
    foo f = new foo();
    Vector<@KeyFor("f.m") String> coll_local = new Vector<>();
    f.m.put(k2, new Object());
    coll_local.add(k2);
  }

  /*public void testAddToListInsteadOfMap(List<@KeyFor("#4") String> la, String b, @KeyFor("#4") String c, Map<String, String> a) {
    // Disabled error (assignment)
    List<String> ls1 = la;
    List<@KeyFor("#4") String> ls2 = la;
    ls1.add(b);
    // Disabled error (argument)
    la.add(b);
    ls2.add(c);
    la.add(c);
    @NonNull String astr = a.get(ls2.get(0));
  }*/
}