package org.checkerframework.framework.type.typeannotator;

import com.sun.tools.javac.code.Type.WildcardType;
import java.util.ArrayDeque;
import java.util.Set;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeKind;
import org.checkerframework.framework.type.AnnotatedTypeFactory;
import org.checkerframework.framework.type.AnnotatedTypeMirror;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedDeclaredType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedWildcardType;
import org.checkerframework.javacutil.BugInCF;
import org.checkerframework.javacutil.PluginUtil;
import org.checkerframework.javacutil.TypesUtils;

/**
 * {@link PropagationTypeAnnotator} adds qualifiers to types where the qualifier to add should be
 * transferred from one or more other types.
 *
 * <p>At the moment, the only function PropagationTypeAnnotator provides, is the propagation of
 * generic type parameter annotations to unannotated wildcards with missing bounds annotations.
 *
 * @see
 *     #visitWildcard(org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedWildcardType,
 *     Object)
 *     <p>PropagationTypeAnnotator traverses trees deeply by default.
 */
public class PropagationTypeAnnotator extends TypeAnnotator {

    // The PropagationTypeAnnotator is called recursively via
    // TypeAnnotatorUtil.eraseBoundsThenAnnotate.
    // This flag prevents infinite recursion.
    private boolean pause = false;
    private ArrayDeque<AnnotatedDeclaredType> parents = new ArrayDeque<>();

    public PropagationTypeAnnotator(AnnotatedTypeFactory typeFactory) {
        super(typeFactory);
    }

    @Override
    public void reset() {
        if (!pause) {
            // when the PropagationTypeAnnotator is called recursively we don't
            // want the visit method to reset the list of visited types
            super.reset();
        }
    }

    /*
     * When pause == true, the PropagationTypeAnnotator caused a recursive call
     * and there is no need to execute the PropagationTypeAnnotator
     */
    @Override
    protected Void scan(AnnotatedTypeMirror type, Void aVoid) {
        if (pause) {
            return null;
        }

        return super.scan(type, aVoid);
    }

    /**
     * Sometimes the underlying type parameters of AnnotatedWildcardTypes are not available on the
     * wildcards themselves. Instead, record enclosing class to find the type parameter to use as a
     * backup in visitWildcards.
     *
     * @param declaredType type to record
     */
    @Override
    public Void visitDeclared(AnnotatedDeclaredType declaredType, Void aVoid) {
        if (pause) {
            return null;
        }
        parents.addFirst(declaredType);
        super.visitDeclared(declaredType, aVoid);
        parents.removeFirst();
        return null;
    }

    /**
     * Rather than defaulting the missing bounds of a wildcard, find the bound annotations on the
     * type parameter it replaced. Place those annotations on the wildcard.
     *
     * @param wildcardAtm type to annotate
     */
    @Override
    public Void visitWildcard(AnnotatedWildcardType wildcardAtm, Void aVoid) {
        if (visitedNodes.containsKey(wildcardAtm) || pause) {
            return null;
        }
        visitedNodes.put(wildcardAtm, null);

        final WildcardType wildcard = (WildcardType) wildcardAtm.getUnderlyingType();
        Element typeParamElement = TypesUtils.wildcardToTypeParam(wildcard);
        if (typeParamElement == null) {
            typeParamElement =
                    parents.isEmpty()
                            ? null
                            : getTypeParamFromEnclosingClass(wildcardAtm, parents.peekFirst());
        }

        if (typeParamElement != null) {
            pause = true;
            AnnotatedTypeVariable typeParam =
                    (AnnotatedTypeVariable) typeFactory.getAnnotatedType(typeParamElement);
            pause = false;

            final Set<? extends AnnotationMirror> tops =
                    typeFactory.getQualifierHierarchy().getTopAnnotations();

            if (wildcard.isUnbound()) {
                propagateExtendsBound(wildcardAtm, typeParam, tops);
                propagateSuperBound(wildcardAtm, typeParam, tops);

            } else if (wildcard.isExtendsBound()) {
                propagateSuperBound(wildcardAtm, typeParam, tops);

            } else { // is super bound
                propagateExtendsBound(wildcardAtm, typeParam, tops);
            }
        }
        scan(wildcardAtm.getExtendsBound(), null);
        scan(wildcardAtm.getSuperBound(), null);
        return null;
    }

    private void propagateSuperBound(
            AnnotatedWildcardType wildcard,
            AnnotatedTypeVariable typeParam,
            Set<? extends AnnotationMirror> tops) {
        applyAnnosFromBound(wildcard.getSuperBound(), typeParam.getLowerBound(), tops);
    }

    private void propagateExtendsBound(
            AnnotatedWildcardType wildcard,
            AnnotatedTypeVariable typeParam,
            Set<? extends AnnotationMirror> tops) {
        applyAnnosFromBound(wildcard.getExtendsBound(), typeParam.getUpperBound(), tops);
    }

    /**
     * Take the primary annotations from typeParamBound and place them as primary annotations on
     * wildcard bound.
     */
    private void applyAnnosFromBound(
            final AnnotatedTypeMirror wildcardBound,
            final AnnotatedTypeMirror typeParamBound,
            final Set<? extends AnnotationMirror> tops) {
        // Type variables do not need primary annotations.
        // The type variable will have annotations placed on its
        // bounds via its declaration or defaulting rules
        if (wildcardBound.getKind() == TypeKind.TYPEVAR
                || typeParamBound.getKind() == TypeKind.TYPEVAR) {
            return;
        }

        for (final AnnotationMirror top : tops) {
            if (wildcardBound.getAnnotationInHierarchy(top) == null) {
                final AnnotationMirror typeParamAnno = typeParamBound.getAnnotationInHierarchy(top);
                if (typeParamAnno == null) {
                    throw new BugInCF(
                            PluginUtil.joinLines(
                                    "Missing annotation on type parameter",
                                    "top=" + top,
                                    "wildcardBound=" + wildcardBound,
                                    "typeParamBound=" + typeParamBound));
                } // else
                wildcardBound.addAnnotation(typeParamAnno);
            }
        }
    }

    /**
     * Search parent's type arguments for wildcard. Using the index of wildcard, find the
     * corresponding type parameter element and return it. Returns null if the wildcard is the
     * result of substitution and therefore not in the list of type arguments.
     */
    private Element getTypeParamFromEnclosingClass(
            final AnnotatedWildcardType wildcard, final AnnotatedDeclaredType parent) {
        Integer wildcardIndex = null;
        int currentIndex = 0;
        for (AnnotatedTypeMirror typeArg : parent.getTypeArguments()) {
            // the only cases in which the wildcard is not one of the type arguments are cases in
            // which they should have been replaced by capture
            if (typeArg == wildcard) {
                wildcardIndex = currentIndex;
                break;
            }
            currentIndex += 1;
        }

        if (wildcardIndex != null) {
            final TypeElement typeElement =
                    (TypeElement)
                            typeFactory
                                    .getProcessingEnv()
                                    .getTypeUtils()
                                    .asElement(parent.getUnderlyingType());

            return typeElement.getTypeParameters().get(wildcardIndex);
        }

        return null;
    }
}
