import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

public class TemporaryConcurrentMap<K, V> extends ConcurrentHashMap<K, V> {

    private final long defaultCooldownTime;

    private final ConcurrentHashMap<K, Long> cooldownMap = new ConcurrentHashMap<>();

    public TemporaryConcurrentMap(long defaultCooldown) {
        this.defaultCooldownTime = defaultCooldown;
    }

    @Override
    public V put(K key, V value) {
        return put(key, value, defaultCooldownTime);
    }

    public V put(K key, V value, Long cooldownTime) {
        V out = super.put(key, value);
        if (!Objects.isNull(out)) {
            this.cooldownMap.put(key, System.currentTimeMillis() + cooldownTime);
        }
        return out;
    }

    @Override
    public void putAll(Map<? extends K, ? extends V> m) {
        putAll(m, defaultCooldownTime);
    }

    public void putAll(Map<? extends K, ? extends V> m, Long cooldownTime) {
        m.entrySet().forEach((e) -> {
            this.cooldownMap.put(e.getKey(), System.currentTimeMillis() + cooldownTime);
        });
        super.putAll(m);
    }

    @Override
    public boolean containsKey(Object key) {
        Long cooldown = cooldownMap.getOrDefault(key, 0l);
        if (cooldown >= System.currentTimeMillis()) {
            return super.containsKey(key);
        }
        this.remove(key);
        return false;
    }

    @Override
    public boolean containsValue(Object value) throws UnsupportedOperationException {
        throw new UnsupportedOperationException();
    }

    @Override
    public void clear() {
        super.clear();
        cooldownMap.clear();
    }

    @Override
    public V get(Object key) {
        Long cooldown = cooldownMap.getOrDefault(key, 0l);
        if (cooldown >= System.currentTimeMillis()) {
            return super.get(key);
        }
        this.remove(key);
        return null;
    }

    @Override
    public V getOrDefault(Object key, V defaultValue) {
        V v;
        return (v = this.get(key)) == null ? defaultValue : v;
    }

    @Override
    public V remove(Object key) {
        cooldownMap.remove(key);
        return super.remove(key);
    }

    @Override
    public boolean remove(Object key, Object value) throws UnsupportedOperationException {
        throw new UnsupportedOperationException();
    }
}