package org.checkerframework.framework.type;

import java.util.IdentityHashMap;
import java.util.Map;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;

/** Duplicates annotated types and replaces components according to a replacement map. */
public class AnnotatedTypeReplacer {

    /**
     * Return a copy of type after making the specified replacements.
     *
     * @param type the type that will be copied with replaced components
     * @param replacementMap a mapping of {@literal referenceToReplace ⇒
     *     referenceOfReplacement}
     * @return a duplicate of type in which every reference that was a key in replacementMap has
     *     been replaced by its corresponding value
     */
    public static AnnotatedTypeMirror replace(
            AnnotatedTypeMirror type,
            Map<? extends AnnotatedTypeMirror, ? extends AnnotatedTypeMirror> replacementMap) {
        return new Visitor(replacementMap).visit(type);
    }

    /**
     * AnnotatedTypeCopier maintains a mapping of {@literal typeVisited ⇒ copyOfTypeVisited}
     * When a reference, typeVisited, is encountered again, it will use the recorded reference,
     * copyOfTypeVisited, instead of generating a new copy of typeVisited. Visitor pre-populates
     * this mapping so that references are replaced not by their copies but by those in the
     * replacementMap provided in the constructor.
     *
     * <p>All types NOT in the replacement map are duplicated as per AnnotatedTypeCopier.visit
     */
    protected static class Visitor extends AnnotatedTypeCopier {

        private final IdentityHashMap<? extends AnnotatedTypeMirror, ? extends AnnotatedTypeMirror>
                originalMappings;

        public Visitor(
                final Map<? extends AnnotatedTypeMirror, ? extends AnnotatedTypeMirror> mappings) {
            originalMappings = new IdentityHashMap<>(mappings);
        }

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

        @Override
        public AnnotatedTypeMirror visitTypeVariable(
                AnnotatedTypeVariable original,
                IdentityHashMap<AnnotatedTypeMirror, AnnotatedTypeMirror> originalToCopy) {
            // AnnotatedTypeCopier will visit the type parameters of a method and copy them
            // Without this flag, any mappings in originalToCopy would replace the type parameters.
            // However, we do not replace the type parameters in an AnnotatedExecutableType.  Also,
            // AnnotatedExecutableType.typeVarTypes is of type List<AnnotatedTypeVariable> so if
            // the mapping contained a type parameter -> (Non-type variable AnnotatedTypeMirror) then
            // a runtime exception would occur.
            if (visitingExecutableTypeParam) {
                visitingExecutableTypeParam = false;
                final AnnotatedTypeVariable copy =
                        (AnnotatedTypeVariable)
                                AnnotatedTypeMirror.createType(
                                        original.getUnderlyingType(),
                                        original.atypeFactory,
                                        original.isDeclaration());
                maybeCopyPrimaryAnnotations(original, copy);
                originalToCopy.put(original, copy);

                if (original.getUpperBoundField() != null) {
                    copy.setUpperBound(visit(original.getUpperBoundField(), originalToCopy));
                }

                if (original.getLowerBoundField() != null) {
                    copy.setLowerBound(visit(original.getLowerBoundField(), originalToCopy));
                }
                return copy;
            }

            return super.visitTypeVariable(original, originalToCopy);
        }
    }
}