package org.checkerframework.checker.nullness;

import com.sun.source.tree.ExpressionTree;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import org.checkerframework.checker.nullness.qual.KeyFor;
import org.checkerframework.framework.flow.CFAbstractAnalysis;
import org.checkerframework.framework.flow.CFAbstractValue;
import org.checkerframework.javacutil.AnnotationUtils;

/**
 * KeyForValue holds additional information about which maps this value is a key for. This extra
 * information is required when adding the @KeyFor qualifier to the type is not a refinement of the
 * type. For example,
 *
 * <pre>
 * {@code Map<T, Object> map = ...;}
 * {@code <T> T method(T param) { }
 *   if (map.contains(param) {
 *     {@code @NonNull Object o = map.get(param);}
 *     return param;
 *   }
 * }
 * }</pre>
 *
 * Inside the if statement, {@code param} is a key for "map". This would normally be represented as
 * {@code @KeyFor("map") T}, but this is not a subtype of {@code T}, so the type cannot be refined.
 * Instead, the value for {@code param} includes "map" in the list of keyForMaps. This information
 * is used in {@link KeyForAnnotatedTypeFactory#isKeyForMap(String, ExpressionTree)}.
 */
public class KeyForValue extends CFAbstractValue<KeyForValue> {
    /**
     * If the underlying type is a type variable or a wildcard, then this is a set of maps for which
     * this value is a key. Otherwise, it's null.
     */
    // Cannot be final because lub re-assigns; add a new constructor to do this cleanly?
    private Set<String> keyForMaps;

    /** Create an instance. */
    public KeyForValue(
            CFAbstractAnalysis<KeyForValue, ?, ?> analysis,
            Set<AnnotationMirror> annotations,
            TypeMirror underlyingType) {
        super(analysis, annotations, underlyingType);
        AnnotationMirror keyfor = AnnotationUtils.getAnnotationByClass(annotations, KeyFor.class);
        if (keyfor != null
                && (underlyingType.getKind() == TypeKind.TYPEVAR
                        || underlyingType.getKind() == TypeKind.WILDCARD)) {
            keyForMaps = new LinkedHashSet<>();
            List<String> list =
                    AnnotationUtils.getElementValueArray(keyfor, "value", String.class, true);
            keyForMaps.addAll(list);
        } else {
            keyForMaps = null;
        }
    }

    /**
     * If the underlying type is a type variable or a wildcard, then this is a set of maps for which
     * this value is a key. Otherwise, it's null.
     */
    public Set<String> getKeyForMaps() {
        return keyForMaps;
    }

    @Override
    public KeyForValue leastUpperBound(KeyForValue other) {
        KeyForValue lub = super.leastUpperBound(other);
        if (other == null || other.keyForMaps == null || this.keyForMaps == null) {
            return lub;
        }
        // Lub the keyForMaps by intersecting the sets.
        lub.keyForMaps = new LinkedHashSet<>();
        lub.keyForMaps.addAll(this.keyForMaps);
        lub.keyForMaps.retainAll(other.keyForMaps);
        if (lub.keyForMaps.isEmpty()) {
            lub.keyForMaps = null;
        }
        return lub;
    }

    @Override
    public KeyForValue mostSpecific(KeyForValue other, KeyForValue backup) {
        KeyForValue mostSpecific = super.mostSpecific(other, backup);
        if (mostSpecific == null) {
            if (other == null) {
                return this;
            }
            // mostSpecific is null if the two types are not comparable.  This is normally
            // because one of this or other is a type variable and annotations is empty, but the
            // other annotations are not empty.  In this case, copy the keyForMaps and to the
            // value with the no annotations and return it as most specific.
            if (other.getAnnotations().isEmpty()) {
                other.addKeyFor(this.keyForMaps);
                return other;
            } else if (this.getAnnotations().isEmpty()) {
                this.addKeyFor(other.keyForMaps);
                return this;
            }
            return null;
        }

        mostSpecific.addKeyFor(this.keyForMaps);
        if (other != null) {
            mostSpecific.addKeyFor(other.keyForMaps);
        }
        return mostSpecific;
    }

    private void addKeyFor(Set<String> newKeyForMaps) {
        if (newKeyForMaps == null || newKeyForMaps.isEmpty()) {
            return;
        }
        if (keyForMaps == null) {
            keyForMaps = new LinkedHashSet<>();
        }
        keyForMaps.addAll(newKeyForMaps);
    }
}
