package org.checkerframework.framework.type;

import java.util.ArrayList;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.List;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedArrayType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedDeclaredType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedExecutableType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedIntersectionType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedNoType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedNullType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedPrimitiveType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedUnionType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedWildcardType;
import org.checkerframework.framework.type.visitor.AnnotatedTypeVisitor;

/**
 * AnnotatedTypeCopier is a visitor that deep copies an AnnotatedTypeMirror exactly, including any
 * lazily initialized fields. That is, if a field has already been initialized, it will be
 * initialized in the copied type.
 *
 * <p>When making copies, a map of encountered {@literal references &rArr; copied} types is
 * maintained. This ensures that, if a reference appears in multiple locations in the original type,
 * a corresponding copy of the original type appears in the same locations in the output copy. This
 * ensures that the recursive loops in the input type are preserved in its output copy (see
 * makeOrReturnCopy)
 *
 * <p>In general, AnnotatedTypeMirrors should be copied via AnnotatedTypeMirror#deepCopy and
 * AnnotatedTypeMirror#shallowCopy. AnnotatedTypeMirror#deepCopy makes use of AnnotatedTypeCopier
 * under the covers. However, this visitor and its subclasses can be invoked as follows:
 *
 * <pre>{@code new AnnotatedTypeCopier().visit(myTypeVar);}</pre>
 *
 * Note: There are methods that may require a copy of a type mirror with slight changes. It is
 * intended that this class can be overridden for these cases.
 *
 * @see org.checkerframework.framework.type.TypeVariableSubstitutor
 * @see org.checkerframework.framework.type.AnnotatedTypeReplacer
 */
public class AnnotatedTypeCopier
        implements AnnotatedTypeVisitor<
                AnnotatedTypeMirror, IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror>> {

    /**
     * This is a hack to handle the curious behavior of substitution on an AnnotatedExecutableType.
     *
     * @see org.checkerframework.framework.type.TypeVariableSubstitutor It is poor form to include
     *     such a flag on the base class for exclusive use in a subclass but it is the least bad
     *     option in this case.
     */
    protected boolean visitingExecutableTypeParam = false;

    /** @see #AnnotatedTypeCopier(boolean) */
    protected final boolean copyAnnotations;

    /**
     * Creates an AnnotatedTypeCopier that may or may not copyAnnotations By default
     * AnnotatedTypeCopier provides two major properties in its copies:
     *
     * <ol>
     *   <li>Structure preservation -- the exact structure of the original AnnotatedTypeMirror is
     *       preserved in the copy including all component types.
     *   <li>Annotation preservation -- All of the annotations from the original AnnotatedTypeMirror
     *       and its components have been copied to the new type.
     * </ol>
     *
     * If copyAnnotations is set to false, the second property, annotation preservation, is removed.
     * This is useful for cases in which the user may want to copy the structure of a type exactly
     * but NOT its annotations.
     */
    public AnnotatedTypeCopier(final boolean copyAnnotations) {
        this.copyAnnotations = copyAnnotations;
    }

    /**
     * Creates an AnnotatedTypeCopier that copies both the structure and annotations of the source
     * AnnotatedTypeMirror.
     *
     * @see #AnnotatedTypeCopier(boolean)
     */
    public AnnotatedTypeCopier() {
        this(true);
    }

    @Override
    public AnnotatedTypeMirror visit(AnnotatedTypeMirror type) {
        return type.accept(this, new IdentityHashMap<>());
    }

    @Override
    public AnnotatedTypeMirror visit(
            AnnotatedTypeMirror type,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        return type.accept(this, originalToCopy);
    }

    @Override
    public AnnotatedTypeMirror visitDeclared(
            AnnotatedDeclaredType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedDeclaredType copy =
                (AnnotatedDeclaredType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        if (original.wasRaw()) {
            copy.setWasRaw();
        }

        if (original.enclosingType != null) {
            copy.enclosingType =
                    (AnnotatedDeclaredType) visit(original.enclosingType, originalToCopy);
        }

        if (original.typeArgs != null) {
            final List<AnnotatedTypeMirror> copyTypeArgs = new ArrayList<>();
            for (final AnnotatedTypeMirror typeArg : original.getTypeArguments()) {
                copyTypeArgs.add(visit(typeArg, originalToCopy));
            }
            copy.setTypeArguments(copyTypeArgs);
        }

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitIntersection(
            AnnotatedIntersectionType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedIntersectionType copy =
                (AnnotatedIntersectionType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        if (original.supertypes != null) {
            final List<AnnotatedDeclaredType> copySupertypes = new ArrayList<>();
            for (final AnnotatedDeclaredType supertype : original.supertypes) {
                copySupertypes.add((AnnotatedDeclaredType) visit(supertype, originalToCopy));
            }
            copy.supertypes = Collections.unmodifiableList(copySupertypes);
        }

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitUnion(
            AnnotatedUnionType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedUnionType copy =
                (AnnotatedUnionType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        if (original.alternatives != null) {
            final List<AnnotatedDeclaredType> copyAlternatives = new ArrayList<>();
            for (final AnnotatedDeclaredType supertype : original.alternatives) {
                copyAlternatives.add((AnnotatedDeclaredType) visit(supertype, originalToCopy));
            }
            copy.alternatives = Collections.unmodifiableList(copyAlternatives);
        }

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitExecutable(
            AnnotatedExecutableType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedExecutableType copy =
                (AnnotatedExecutableType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        copy.setElement(original.getElement());

        if (original.receiverType != null) {
            copy.receiverType =
                    (AnnotatedDeclaredType) visit(original.receiverType, originalToCopy);
        }

        for (final AnnotatedTypeMirror param : original.paramTypes) {
            copy.paramTypes.add(visit(param, originalToCopy));
        }

        for (final AnnotatedTypeMirror thrown : original.throwsTypes) {
            copy.throwsTypes.add(visit(thrown, originalToCopy));
        }

        copy.returnType = visit(original.returnType, originalToCopy);

        for (final AnnotatedTypeVariable typeVariable : original.typeVarTypes) {
            // This field is needed to identify exactly when the declaration of an executable's
            // type parameter is visited.  When subtypes of this class visit the type parameter's
            // component types, they will likely set visitingExecutableTypeParam to false.
            // Therefore, we set this variable on each iteration of the loop.
            // See TypeVariableSubstitutor.Visitor.visitTypeVariable for an example of this.
            visitingExecutableTypeParam = true;
            copy.typeVarTypes.add((AnnotatedTypeVariable) visit(typeVariable, originalToCopy));
        }
        visitingExecutableTypeParam = false;

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitArray(
            AnnotatedArrayType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedArrayType copy =
                (AnnotatedArrayType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        copy.setComponentType(visit(original.getComponentType(), originalToCopy));

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitTypeVariable(
            AnnotatedTypeVariable original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedTypeVariable copy =
                (AnnotatedTypeVariable)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        if (original.getUpperBoundField() != null) {
            // TODO: figure out why asUse is needed here and remove it.
            copy.setUpperBound(visit(original.getUpperBoundField(), originalToCopy).asUse());
        }

        if (original.getLowerBoundField() != null) {
            // TODO: figure out why asUse is needed here and remove it.
            copy.setLowerBound(visit(original.getLowerBoundField(), originalToCopy).asUse());
        }

        return copy;
    }

    @Override
    public AnnotatedTypeMirror visitPrimitive(
            AnnotatedPrimitiveType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        return makeCopy(original);
    }

    @Override
    public AnnotatedTypeMirror visitNoType(
            AnnotatedNoType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        return makeCopy(original);
    }

    @Override
    public AnnotatedTypeMirror visitNull(
            AnnotatedNullType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        return makeOrReturnCopy(original, originalToCopy);
    }

    @Override
    public AnnotatedTypeMirror visitWildcard(
            AnnotatedWildcardType original,
            IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return originalToCopy.get(original);
        }

        final AnnotatedWildcardType copy =
                (AnnotatedWildcardType)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        if (original.isUninferredTypeArgument()) {
            copy.setUninferredTypeArgument();
        }

        maybeCopyPrimaryAnnotations(original, copy);
        originalToCopy.put(original, copy);

        if (original.getExtendsBoundField() != null) {
            copy.setExtendsBound(visit(original.getExtendsBoundField(), originalToCopy).asUse());
        }

        if (original.getSuperBoundField() != null) {
            copy.setSuperBound(visit(original.getSuperBoundField(), originalToCopy).asUse());
        }

        copy.setTypeVariable(original.getTypeVariable());

        return copy;
    }

    /**
     * For any given object in the type being copied, we only want to generate one copy of that
     * object. When that object is encountered again, using the previously generated copy will
     * preserve the structure of the original AnnotatedTypeMirror.
     *
     * <p>makeOrReturnCopy first checks to see if an object has been encountered before. If so, it
     * returns the previously generated duplicate of that object if not, it creates a duplicate of
     * the object and stores it in the history, originalToCopy
     *
     * @param original a reference to a type to copy
     * @param originalToCopy a mapping of previously encountered references to the copies made for
     *     those references
     * @param <T> the type of original copy, this is a shortcut to avoid having to insert casts all
     *     over the visitor
     * @return a copy of original
     */
    @SuppressWarnings("unchecked")
    protected <T extends AnnotatedTypeMirror> T makeOrReturnCopy(
            T original, IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
        if (originalToCopy.containsKey(original)) {
            return (T) originalToCopy.get(original);
        }

        final T copy = makeCopy(original);
        originalToCopy.put(original, copy);

        return copy;
    }

    @SuppressWarnings("unchecked")
    protected <T extends AnnotatedTypeMirror> T makeCopy(T original) {

        final T copy =
                (T)
                        AnnotatedTypeMirror.createType(
                                original.getUnderlyingType(),
                                original.atypeFactory,
                                original.isDeclaration());
        maybeCopyPrimaryAnnotations(original, copy);

        return copy;
    }

    /**
     * This method is called in any location in which a primary annotation would be copied from
     * source to dest. Note, this method obeys the copyAnnotations field. Subclasses of
     * AnnotatedTypeCopier can use this method to customize annotations before copying.
     *
     * @param source the type whose primary annotations are being copied
     * @param dest a copy of source that should receive its primary annotations
     */
    protected void maybeCopyPrimaryAnnotations(
            final AnnotatedTypeMirror source, final AnnotatedTypeMirror dest) {
        if (copyAnnotations) {
            dest.addAnnotations(source.getAnnotationsField());
        }
    }
}
