package org.checkerframework.framework.type;

import java.util.Set;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.type.TypeVariable;
import javax.lang.model.type.WildcardType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedIntersectionType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;
import org.checkerframework.framework.util.AnnotatedTypes;
import org.checkerframework.javacutil.BugInCF;

/** Utility class for applying the annotations inferred by dataflow to a given type. */
public class DefaultInferredTypesApplier {

    // At the moment, only Inference uses the omitSubtypingCheck option.
    // In actuality the subtyping check should be unnecessary since inferred
    // types should be subtypes of their declaration.
    private final boolean omitSubtypingCheck;

    private final QualifierHierarchy hierarchy;
    private final AnnotatedTypeFactory factory;

    public DefaultInferredTypesApplier(QualifierHierarchy hierarchy, AnnotatedTypeFactory factory) {
        this(false, hierarchy, factory);
    }

    public DefaultInferredTypesApplier(
            boolean omitSubtypingCheck,
            QualifierHierarchy hierarchy,
            AnnotatedTypeFactory factory) {
        this.omitSubtypingCheck = omitSubtypingCheck;
        this.hierarchy = hierarchy;
        this.factory = factory;
    }

    /**
     * For each top in qualifier hierarchy, traverse inferred and copy the required annotations over
     * to type.
     *
     * @param type the type to which annotations are being applied
     * @param inferredSet the type inferred by data flow
     * @param inferredTypeMirror underlying inferred type
     */
    public void applyInferredType(
            final AnnotatedTypeMirror type,
            final Set<AnnotationMirror> inferredSet,
            TypeMirror inferredTypeMirror) {
        if (inferredSet == null) {
            return;
        }
        if (inferredTypeMirror.getKind() == TypeKind.WILDCARD) {
            // Dataflow might infer a wildcard that extends a type variable for types that are
            // actually type variables.  Use the type variable instead.
            while (inferredTypeMirror.getKind() == TypeKind.WILDCARD
                    && (((WildcardType) inferredTypeMirror).getExtendsBound() != null)) {
                inferredTypeMirror = ((WildcardType) inferredTypeMirror).getExtendsBound();
            }
        }
        for (final AnnotationMirror top : hierarchy.getTopAnnotations()) {
            AnnotationMirror inferred = hierarchy.findAnnotationInHierarchy(inferredSet, top);

            apply(type, inferred, inferredTypeMirror, top);
        }
    }

    private void apply(
            AnnotatedTypeMirror type,
            AnnotationMirror inferred,
            TypeMirror inferredTypeMirror,
            AnnotationMirror top) {
        AnnotationMirror primary = type.getAnnotationInHierarchy(top);
        if (inferred == null) {

            if (primary == null) {
                // Type doesn't have a primary either, nothing to remove
            } else if (type.getKind() == TypeKind.TYPEVAR) {
                removePrimaryAnnotationTypeVar(
                        (AnnotatedTypeVariable) type, inferredTypeMirror, top, primary);
            } else {
                removePrimaryTypeVarApplyUpperBound(type, inferredTypeMirror, top, primary);
            }
        } else {
            if (primary == null) {
                Set<AnnotationMirror> lowerbounds =
                        AnnotatedTypes.findEffectiveLowerBoundAnnotations(hierarchy, type);
                primary = hierarchy.findAnnotationInHierarchy(lowerbounds, top);
            }
            if ((omitSubtypingCheck || hierarchy.isSubtype(inferred, primary))) {
                type.replaceAnnotation(inferred);
                if (type.getKind() == TypeKind.INTERSECTION) {
                    for (AnnotatedTypeMirror superType :
                            ((AnnotatedIntersectionType) type).directSuperTypes()) {
                        superType.replaceAnnotation(inferred);
                    }
                }
            }
        }
    }

    private void removePrimaryTypeVarApplyUpperBound(
            AnnotatedTypeMirror type,
            TypeMirror inferredTypeMirror,
            AnnotationMirror top,
            AnnotationMirror notInferred) {
        if (inferredTypeMirror.getKind() != TypeKind.TYPEVAR) {
            throw new BugInCF(
                    "Inferred value should not be missing annotations: " + inferredTypeMirror);
        }

        TypeVariable typeVar = (TypeVariable) inferredTypeMirror;
        AnnotatedTypeVariable typeVariableDecl =
                (AnnotatedTypeVariable) factory.getAnnotatedType(typeVar.asElement());
        AnnotationMirror upperBound = typeVariableDecl.getEffectiveAnnotationInHierarchy(top);

        if (omitSubtypingCheck || hierarchy.isSubtype(upperBound, notInferred)) {
            type.replaceAnnotation(upperBound);
        }
    }

    private void removePrimaryAnnotationTypeVar(
            AnnotatedTypeVariable annotatedTypeVariable,
            TypeMirror inferredTypeMirror,
            AnnotationMirror top,
            AnnotationMirror previousAnnotation) {
        if (inferredTypeMirror.getKind() != TypeKind.TYPEVAR) {
            throw new BugInCF("Missing annos");
        }
        TypeVariable typeVar = (TypeVariable) inferredTypeMirror;
        AnnotatedTypeVariable typeVariableDecl =
                (AnnotatedTypeVariable) factory.getAnnotatedType(typeVar.asElement());
        AnnotationMirror upperBound = typeVariableDecl.getEffectiveAnnotationInHierarchy(top);
        if (omitSubtypingCheck || hierarchy.isSubtype(upperBound, previousAnnotation)) {
            annotatedTypeVariable.removeAnnotationInHierarchy(top);
            AnnotationMirror ub = typeVariableDecl.getUpperBound().getAnnotationInHierarchy(top);
            apply(annotatedTypeVariable.getUpperBound(), ub, typeVar.getUpperBound(), top);
            AnnotationMirror lb = typeVariableDecl.getLowerBound().getAnnotationInHierarchy(top);
            apply(annotatedTypeVariable.getLowerBound(), lb, typeVar.getLowerBound(), top);
        }
    }
}
