package org.checkerframework.framework.util.typeinference.solver;

import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeVariable;
import org.checkerframework.framework.type.AnnotatedTypeFactory;
import org.checkerframework.framework.type.AnnotatedTypeMirror;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;
import org.checkerframework.framework.util.AnnotationMirrorMap;
import org.checkerframework.framework.util.AnnotationMirrorSet;
import org.checkerframework.framework.util.typeinference.TypeArgInferenceUtil;
import org.checkerframework.framework.util.typeinference.solver.InferredValue.InferredTarget;
import org.checkerframework.framework.util.typeinference.solver.InferredValue.InferredType;
import org.checkerframework.framework.util.typeinference.solver.TargetConstraints.Equalities;
import org.checkerframework.javacutil.BugInCF;

/**
 * EqualitiesSolver infers type arguments for targets using the equality constraints in
 * ConstraintMap. When a type is inferred, it rewrites the remaining equality/supertype constraints
 */
public class EqualitiesSolver {
    private boolean dirty = false;

    /**
     * For each target, if there is one or more equality constraints involving concrete types that
     * lets us infer a primary annotation in all qualifier hierarchies then infer a concrete type
     * argument. else if there is one or more equality constraints involving other targets that lets
     * us infer a primary annotation in all qualifier hierarchies then infer that type argument is
     * the other type argument
     *
     * <p>if we have inferred either a concrete type or another target as type argument rewrite all
     * of the constraints for the current target to instead use the inferred type/target
     *
     * <p>We do this iteratively until NO new inferred type argument is found
     *
     * @param targets the list of type parameters for which we are inferring type arguments
     * @param constraintMap the set of constraints over the set of targets
     * @return a Map( {@code target &rArr; inferred type or target })
     */
    public InferenceResult solveEqualities(
            Set<TypeVariable> targets,
            ConstraintMap constraintMap,
            AnnotatedTypeFactory typeFactory) {
        final InferenceResult solution = new InferenceResult();

        do {
            dirty = false;
            for (TypeVariable target : targets) {
                if (solution.containsKey(target)) {
                    continue;
                }

                Equalities equalities = constraintMap.getConstraints(target).equalities;

                InferredValue inferred =
                        mergeConstraints(target, equalities, solution, constraintMap, typeFactory);
                if (inferred != null) {
                    if (inferred instanceof InferredType) {
                        rewriteWithInferredType(
                                target, ((InferredType) inferred).type, constraintMap);
                    } else {
                        rewriteWithInferredTarget(
                                target,
                                ((InferredTarget) inferred).target,
                                constraintMap,
                                typeFactory);
                    }

                    solution.put(target, inferred);
                }
            }

        } while (dirty);

        solution.resolveChainedTargets();

        return solution;
    }

    /**
     * Let Ti be a target type parameter. When we reach this method we have inferred an argument,
     * Ai, for Ti
     *
     * <p>However, there still may be constraints of the form {@literal Ti = Tj}, {@literal Ti <:
     * Tj}, {@literal Tj <: Ti} in the constraint map. In this case we need to replace Ti with the
     * type. That is, they become {@literal Ai = Tj}, {@literal Ai <: Tj}, and {@literal Tj <: Ai}
     *
     * <p>To do this, we find the TargetConstraints for Tj and add these constraints to the
     * appropriate map in TargetConstraints. We can then clear the constraints for the current
     * target since we have inferred a type.
     *
     * @param target the target for which we have inferred a concrete type argument
     * @param type the type inferred
     */
    private void rewriteWithInferredType(
            final TypeVariable target,
            final AnnotatedTypeMirror type,
            final ConstraintMap constraints) {

        final TargetConstraints targetRecord = constraints.getConstraints(target);
        final Map<TypeVariable, AnnotationMirrorSet> equivalentTargets =
                targetRecord.equalities.targets;
        // each target that was equivalent to this one needs to be equivalent in the same
        // hierarchies as the inferred type
        for (final Entry<TypeVariable, AnnotationMirrorSet> eqEntry :
                equivalentTargets.entrySet()) {
            constraints.addTypeEqualities(eqEntry.getKey(), type, eqEntry.getValue());
        }

        for (TypeVariable otherTarget : constraints.getTargets()) {
            if (otherTarget != target) {
                final TargetConstraints record = constraints.getConstraints(otherTarget);

                // each target that was equivalent to this one needs to be equivalent in the same
                // hierarchies as the inferred type
                final AnnotationMirrorSet hierarchies = record.equalities.targets.get(target);
                if (hierarchies != null) {
                    record.equalities.targets.remove(target);
                    constraints.addTypeEqualities(otherTarget, type, hierarchies);
                }

                // otherTypes may have AnnotatedTypeVariables of type target, run substitution on
                // these with type
                Map<AnnotatedTypeMirror, AnnotationMirrorSet> toIterate =
                        new LinkedHashMap<>(record.equalities.types);
                record.equalities.types.clear();
                for (AnnotatedTypeMirror otherType : toIterate.keySet()) {
                    final AnnotatedTypeMirror copy =
                            TypeArgInferenceUtil.substitute(target, type, otherType);
                    final AnnotationMirrorSet otherHierarchies = toIterate.get(otherType);
                    record.equalities.types.put(copy, otherHierarchies);
                }
            }
        }

        for (TypeVariable otherTarget : constraints.getTargets()) {
            if (otherTarget != target) {
                final TargetConstraints record = constraints.getConstraints(otherTarget);

                // each target that was equivalent to this one needs to be equivalent in the same
                // hierarchies as the inferred type
                final AnnotationMirrorSet hierarchies = record.supertypes.targets.get(target);
                if (hierarchies != null) {
                    record.supertypes.targets.remove(target);
                    constraints.addTypeEqualities(otherTarget, type, hierarchies);
                }

                // otherTypes may have AnnotatedTypeVariables of type target, run substitution on
                // these with type
                Map<AnnotatedTypeMirror, AnnotationMirrorSet> toIterate =
                        new LinkedHashMap<>(record.supertypes.types);
                record.supertypes.types.clear();
                for (AnnotatedTypeMirror otherType : toIterate.keySet()) {
                    final AnnotatedTypeMirror copy =
                            TypeArgInferenceUtil.substitute(target, type, otherType);
                    final AnnotationMirrorSet otherHierarchies = toIterate.get(otherType);
                    record.supertypes.types.put(copy, otherHierarchies);
                }
            }
        }

        targetRecord.equalities.clear();
        targetRecord.supertypes.clear();
    }

    /**
     * Let Ti be a target type parameter. When we reach this method we have inferred that Ti has the
     * exact same argument as another target Tj
     *
     * <p>Therefore, we want to stop solving for Ti and instead wait till we solve for Tj and use
     * that result.
     *
     * <p>Let ATM be any annotated type mirror and Tk be a target type parameter where k != i and k
     * != j Even though we've inferred Ti = Tj, there still may be constraints of the form Ti = ATM
     * or {@literal Ti <: Tk} These constraints are still useful for inferring a argument for Ti/Tj.
     * So, we replace Ti in these constraints with Tj and place those constraints in the
     * TargetConstraints object for Tj.
     *
     * <p>We then clear the constraints for Ti.
     *
     * @param target the target for which we know another target is exactly equal to this target
     * @param inferredTarget the other target inferred to be equal
     */
    private void rewriteWithInferredTarget(
            final TypeVariable target,
            final TypeVariable inferredTarget,
            final ConstraintMap constraints,
            final AnnotatedTypeFactory typeFactory) {
        final TargetConstraints targetRecord = constraints.getConstraints(target);
        final Map<AnnotatedTypeMirror, AnnotationMirrorSet> equivalentTypes =
                targetRecord.equalities.types;
        final Map<AnnotatedTypeMirror, AnnotationMirrorSet> supertypes =
                targetRecord.supertypes.types;

        // each type that was equivalent to this one needs to be equivalent in the same hierarchies
        // to the inferred target
        for (final Entry<AnnotatedTypeMirror, AnnotationMirrorSet> eqEntry :
                equivalentTypes.entrySet()) {
            constraints.addTypeEqualities(inferredTarget, eqEntry.getKey(), eqEntry.getValue());
        }

        for (final Entry<AnnotatedTypeMirror, AnnotationMirrorSet> superEntry :
                supertypes.entrySet()) {
            constraints.addTypeSupertype(
                    inferredTarget, superEntry.getKey(), superEntry.getValue());
        }

        for (TypeVariable otherTarget : constraints.getTargets()) {
            if (otherTarget != target && otherTarget != inferredTarget) {
                final TargetConstraints record = constraints.getConstraints(otherTarget);

                // each target that was equivalent to this one needs to be equivalent in the same
                // hierarchies as the inferred target
                final AnnotationMirrorSet hierarchies = record.equalities.targets.get(target);
                if (hierarchies != null) {
                    record.equalities.targets.remove(target);
                    constraints.addTargetEquality(otherTarget, inferredTarget, hierarchies);
                }

                // otherTypes may have AnnotatedTypeVariables of type target, run substitution on
                // these with type
                Map<AnnotatedTypeMirror, AnnotationMirrorSet> toIterate =
                        new LinkedHashMap<>(record.equalities.types);
                record.equalities.types.clear();
                for (AnnotatedTypeMirror otherType : toIterate.keySet()) {
                    final AnnotatedTypeMirror copy =
                            TypeArgInferenceUtil.substitute(
                                    target, createAnnotatedTypeVar(target, typeFactory), otherType);
                    final AnnotationMirrorSet otherHierarchies = toIterate.get(otherType);
                    record.equalities.types.put(copy, otherHierarchies);
                }
            }
        }

        for (TypeVariable otherTarget : constraints.getTargets()) {
            if (otherTarget != target && otherTarget != inferredTarget) {
                final TargetConstraints record = constraints.getConstraints(otherTarget);

                final AnnotationMirrorSet hierarchies = record.supertypes.targets.get(target);
                if (hierarchies != null) {
                    record.supertypes.targets.remove(target);
                    constraints.addTargetSupertype(otherTarget, inferredTarget, hierarchies);
                }

                // otherTypes may have AnnotatedTypeVariables of type target, run substitution on
                // these with type
                Map<AnnotatedTypeMirror, AnnotationMirrorSet> toIterate =
                        new LinkedHashMap<>(record.supertypes.types);
                record.supertypes.types.clear();
                for (AnnotatedTypeMirror otherType : toIterate.keySet()) {
                    final AnnotatedTypeMirror copy =
                            TypeArgInferenceUtil.substitute(
                                    target, createAnnotatedTypeVar(target, typeFactory), otherType);
                    final AnnotationMirrorSet otherHierarchies = toIterate.get(otherType);
                    record.supertypes.types.put(copy, otherHierarchies);
                }
            }
        }

        targetRecord.equalities.clear();
        targetRecord.supertypes.clear();
    }

    /** Creates a declaration AnnotatedTypeVariable for TypeVariable. */
    private AnnotatedTypeVariable createAnnotatedTypeVar(
            final TypeVariable typeVariable, final AnnotatedTypeFactory typeFactory) {
        return (AnnotatedTypeVariable) typeFactory.getAnnotatedType(typeVariable.asElement());
    }

    /**
     * @param typesToHierarchies a mapping of (types &rarr; hierarchies) that indicate that the
     *     argument being inferred is equal to the types in each of the hierarchies
     * @param primaries a map (hierarchy &rarr; annotation in hierarchy) where the annotation in
     *     hierarchy is equal to the primary annotation on the argument being inferred
     * @param tops the set of top annotations in the qualifier hierarchy
     * @return a concrete type argument or null if there was not enough information to infer one
     */
    private InferredType mergeTypesAndPrimaries(
            Map<AnnotatedTypeMirror, AnnotationMirrorSet> typesToHierarchies,
            AnnotationMirrorMap<AnnotationMirror> primaries,
            final AnnotationMirrorSet tops,
            AnnotatedTypeFactory typeFactory) {
        final AnnotationMirrorSet missingAnnos = new AnnotationMirrorSet(tops);

        Iterator<Entry<AnnotatedTypeMirror, AnnotationMirrorSet>> entryIterator =
                typesToHierarchies.entrySet().iterator();
        if (!entryIterator.hasNext()) {
            throw new BugInCF("Merging a list of empty types.");
        }

        final Entry<AnnotatedTypeMirror, AnnotationMirrorSet> head = entryIterator.next();

        AnnotatedTypeMirror mergedType = head.getKey();
        missingAnnos.removeAll(head.getValue());

        // 1. if there are multiple equality constraints in a ConstraintMap then the types better
        // have the same underlying type or Javac will complain and we won't be here.  When building
        // ConstraintMaps constraints involving AnnotatedTypeMirrors that are exactly equal are
        // combined so there must be some difference between two types being merged here.
        //
        // 2. Otherwise, we might have the same underlying type but conflicting annotations, then we
        // take the first set of annotations and show an error for the argument/return type that
        // caused the second differing constraint.
        //
        // 3. Finally, we expect the following types to be involved in equality constraints:
        // AnnotatedDeclaredTypes, AnnotatedTypeVariables, and AnnotatedArrayTypes
        while (entryIterator.hasNext() && !missingAnnos.isEmpty()) {
            final Entry<AnnotatedTypeMirror, AnnotationMirrorSet> current = entryIterator.next();
            final AnnotatedTypeMirror currentType = current.getKey();
            final AnnotationMirrorSet currentHierarchies = current.getValue();

            AnnotationMirrorSet found = new AnnotationMirrorSet();
            for (AnnotationMirror top : missingAnnos) {
                if (currentHierarchies.contains(top)) {
                    final AnnotationMirror newAnno = currentType.getAnnotationInHierarchy(top);
                    if (newAnno != null) {
                        mergedType.replaceAnnotation(newAnno);
                        found.add(top);

                    } else if (mergedType.getKind() == TypeKind.TYPEVAR
                            && typeFactory.types.isSameType(
                                    currentType.getUnderlyingType(),
                                    mergedType.getUnderlyingType())) {
                        // the options here are we are merging with the same typevar, in which case
                        // we can just remove the annotation from the missing list
                        found.add(top);

                    } else {
                        // otherwise the other type is missing an annotation
                        throw new BugInCF(
                                "Missing annotation.\n"
                                        + "\nmergedType="
                                        + mergedType
                                        + "\ncurrentType="
                                        + currentType);
                    }
                }
            }

            missingAnnos.removeAll(found);
        }

        // add all the annotations from the primaries
        for (final AnnotationMirror top : missingAnnos) {
            final AnnotationMirror anno = primaries.get(top);
            if (anno != null) {
                mergedType.replaceAnnotation(anno);
            }
        }

        typesToHierarchies.clear();

        if (missingAnnos.isEmpty()) {
            return new InferredType(mergedType);
        }

        // TODO: we probably can do more with this information than just putting it back into the
        // TODO: ConstraintMap (which is what's happening here)
        final AnnotationMirrorSet hierarchies = new AnnotationMirrorSet(tops);
        hierarchies.removeAll(missingAnnos);
        typesToHierarchies.put(mergedType, hierarchies);

        return null;
    }

    public InferredValue mergeConstraints(
            final TypeVariable target,
            final Equalities equalities,
            final InferenceResult solution,
            ConstraintMap constraintMap,
            AnnotatedTypeFactory typeFactory) {
        final AnnotationMirrorSet tops =
                new AnnotationMirrorSet(typeFactory.getQualifierHierarchy().getTopAnnotations());
        InferredValue inferred = null;
        if (!equalities.types.isEmpty()) {
            inferred =
                    mergeTypesAndPrimaries(
                            equalities.types, equalities.primaries, tops, typeFactory);
        }

        if (inferred != null) {
            return inferred;
        } // else

        // We did not have enough information to infer an annotation in all hierarchies for one
        // concrete type.
        // However, we have a "partial solution", one in which we know the type in some but not all
        // qualifier hierarchies.
        // Update our set of constraints with this information
        dirty |= updateTargetsWithPartiallyInferredType(equalities, constraintMap, typeFactory);
        inferred = findEqualTarget(equalities, tops);

        if (inferred == null && equalities.types.size() == 1) {
            // Still could not find an inferred type in all hierarchies, so just use what type is
            // known.
            AnnotatedTypeMirror type = equalities.types.keySet().iterator().next();
            inferred = new InferredType(type);
        }
        return inferred;
    }

    // If we determined that this target T1 is equal to a type ATM in hierarchies @A,@B,@C
    // for each of those hierarchies, if a target is equal to T1 in that hierarchy it is also equal
    // to ATM.
    // e.g.
    //   if : T1 == @A @B @C ATM in only the A,B hierarchies
    //    and T1 == T2 only in @A hierarchy
    //
    //   then T2 == @A @B @C only in the @A hierarchy
    //
    public boolean updateTargetsWithPartiallyInferredType(
            final Equalities equalities,
            ConstraintMap constraintMap,
            AnnotatedTypeFactory typeFactory) {

        boolean updated = false;

        if (!equalities.types.isEmpty()) {
            if (equalities.types.size() != 1) {
                throw new BugInCF("Equalities should have at most 1 constraint.");
            }

            Entry<AnnotatedTypeMirror, AnnotationMirrorSet> remainingTypeEquality;
            remainingTypeEquality = equalities.types.entrySet().iterator().next();

            final AnnotatedTypeMirror remainingType = remainingTypeEquality.getKey();
            final AnnotationMirrorSet remainingHierarchies = remainingTypeEquality.getValue();

            // update targets
            for (Map.Entry<TypeVariable, AnnotationMirrorSet> targetToHierarchies :
                    equalities.targets.entrySet()) {
                final TypeVariable equalTarget = targetToHierarchies.getKey();
                final AnnotationMirrorSet hierarchies = targetToHierarchies.getValue();

                final AnnotationMirrorSet equalTypeHierarchies =
                        new AnnotationMirrorSet(remainingHierarchies);
                equalTypeHierarchies.retainAll(hierarchies);

                final Map<AnnotatedTypeMirror, AnnotationMirrorSet> otherTargetsEqualTypes =
                        constraintMap.getConstraints(equalTarget).equalities.types;

                AnnotationMirrorSet equalHierarchies = otherTargetsEqualTypes.get(remainingType);
                if (equalHierarchies == null) {
                    equalHierarchies = new AnnotationMirrorSet(equalTypeHierarchies);
                    otherTargetsEqualTypes.put(remainingType, equalHierarchies);
                    updated = true;

                } else {
                    final int size = equalHierarchies.size();
                    equalHierarchies.addAll(equalTypeHierarchies);
                    updated = size == equalHierarchies.size();
                }
            }
        }

        return updated;
    }

    /**
     * Attempt to find a target which is equal to this target.
     *
     * @return a target equal to this target in all hierarchies, or null
     */
    public InferredTarget findEqualTarget(final Equalities equalities, AnnotationMirrorSet tops) {
        for (Map.Entry<TypeVariable, AnnotationMirrorSet> targetToHierarchies :
                equalities.targets.entrySet()) {
            final TypeVariable equalTarget = targetToHierarchies.getKey();
            final AnnotationMirrorSet hierarchies = targetToHierarchies.getValue();

            // Now see if target is equal to equalTarget in all hierarchies
            boolean targetIsEqualInAllHierarchies = hierarchies.size() == tops.size();
            if (targetIsEqualInAllHierarchies) {
                return new InferredTarget(equalTarget, new AnnotationMirrorSet());

            } else {
                // annos in primaries that are not covered by the target's list of equal hierarchies
                final AnnotationMirrorSet requiredPrimaries =
                        new AnnotationMirrorSet(equalities.primaries.keySet());
                requiredPrimaries.removeAll(hierarchies);

                boolean typeWithPrimariesIsEqual =
                        (requiredPrimaries.size() + hierarchies.size()) == tops.size();
                if (typeWithPrimariesIsEqual) {
                    return new InferredTarget(equalTarget, requiredPrimaries);
                }
            }
        }

        return null;
    }
}
