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

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.type.QualifierHierarchy;
import org.checkerframework.framework.util.AnnotationMirrorSet;
import org.checkerframework.framework.util.typeinference.constraint.TIsU;
import org.checkerframework.framework.util.typeinference.constraint.TSuperU;
import org.checkerframework.framework.util.typeinference.constraint.TUConstraint;
import org.checkerframework.javacutil.TypeAnnotationUtils;

/** Converts a set of TUConstraints into a ConstraintMap. */
public class ConstraintMapBuilder {

  /**
   * Let Ti be a the ith target being inferred Let ATV(i) be the annotated type variable that
   * represents as use of Ti which may or may not have primary annotations. Let ATM be an annotated
   * type mirror that may or may not be target Tx, or have a component target Tx Let Ai be the type
   * argument we are trying to infer for Ti
   *
   * <p>We have a set of constraints of the form: {@code ATV(i) <?> ATM}
   *
   * <p>Where {@code <?>} is either a subtype ({@code <:}), supertype ({@code :>}), or equality
   * relationship ({@code =}).
   *
   * <p>Regardless of what {@code <?>} is, a constraint will only imply constraints on Ai in a given
   * hierarchy if ATV(i) does NOT have a primary annotation in that hierarchy. That is:
   *
   * <p>E.g. Let ATV(i) be @NonNull Ti, the constraints @NonNull Ti = @NonNull @Initialized String
   * does not imply any primary annotation in the Nullness hierarchy for type argument Ai because
   * the Annotated type mirror has a primary annotation in the NUllness hierarchy.
   *
   * <p>However, it does imply that Ai has a primary annotation of @Initialized since ATV(i) has no
   * primary annotation in the initialization hierarchy.
   *
   * <p>Note, constraints come in 2 forms:
   *
   * <ul>
   *   <li>between a target and a concrete AnnotatedTypeMirror. E.g., As seen above {@code (@NonNull
   *       Ti = @NonNull @Initialized String)}
   *   <li>between two targets E.g., {@code (@NonNull Ti = Tj)}
   * </ul>
   */
  public ConstraintMap build(
      Set<TypeVariable> targets, Set<TUConstraint> constraints, AnnotatedTypeFactory typeFactory) {

    final QualifierHierarchy qualifierHierarchy = typeFactory.getQualifierHierarchy();
    final AnnotationMirrorSet tops =
        new AnnotationMirrorSet(qualifierHierarchy.getTopAnnotations());
    final ConstraintMap result = new ConstraintMap(targets);

    final AnnotationMirrorSet tAnnos = new AnnotationMirrorSet();
    final AnnotationMirrorSet uAnnos = new AnnotationMirrorSet();
    final AnnotationMirrorSet hierarchiesInRelation = new AnnotationMirrorSet();

    for (TUConstraint constraint : constraints) {
      tAnnos.clear();
      uAnnos.clear();
      hierarchiesInRelation.clear();

      final AnnotatedTypeVariable typeT = constraint.typeVariable;
      final AnnotatedTypeMirror typeU = constraint.relatedType;

      // If typeU is from an argument to the method, then treat typeU as an ordinary type even
      // if it is a target type variable.  This is for the case where the inferred type is the
      // declared type parameter.  For example,
      // public <T> T get(T t) {
      //   return this.get(t);
      // }
      // The inferred type of T should be T.
      if (!constraint.uIsArg
          && typeU.getKind() == TypeKind.TYPEVAR
          && targets.contains(
              (TypeVariable) TypeAnnotationUtils.unannotatedType(typeU.getUnderlyingType()))) {
        if (typeT.getAnnotations().isEmpty() && typeU.getAnnotations().isEmpty()) {
          hierarchiesInRelation.addAll(tops);

        } else {

          for (AnnotationMirror top : tops) {
            final AnnotationMirror tAnno = typeT.getAnnotationInHierarchy(top);
            final AnnotationMirror uAnno = typeU.getAnnotationInHierarchy(top);

            if (tAnno == null) {
              if (uAnno == null) {
                hierarchiesInRelation.add(top);

              } else {
                tAnnos.add(uAnno);
              }
            } else {
              if (uAnno == null) {
                uAnnos.add(tAnno);

              } else {
                // This tells us nothing, they both should be equal but either way
                // we gain no information if both type vars have annotations
              }
            }
          }

          // If we have a case where Ti = @NonNull Tj we know that for the @Initialization hierarchy
          // Ti = TJ and we know that for the @Nullable hierarchy Ti = @NonNull <some other type>.
          // This step saves @NonNull annotation.
          // This case also covers the case where i = j.
          if (!tAnnos.isEmpty()) {
            addToPrimaryRelationship(
                (TypeVariable) TypeAnnotationUtils.unannotatedType(typeT.getUnderlyingType()),
                constraint,
                result,
                tAnnos,
                qualifierHierarchy);
          }

          if (!uAnnos.isEmpty()) {
            addToPrimaryRelationship(
                (TypeVariable) TypeAnnotationUtils.unannotatedType(typeU.getUnderlyingType()),
                constraint,
                result,
                uAnnos,
                qualifierHierarchy);
          }
        }

        // This is the case where we have a relationship between two different targets (Ti
        // <?> Tj and i != j)
        if (!typeFactory.types.isSameType(
            TypeAnnotationUtils.unannotatedType(typeT.getUnderlyingType()),
            TypeAnnotationUtils.unannotatedType(typeU.getUnderlyingType()))) {
          addToTargetRelationship(
              (TypeVariable) TypeAnnotationUtils.unannotatedType(typeT.getUnderlyingType()),
              (TypeVariable) TypeAnnotationUtils.unannotatedType(typeU.getUnderlyingType()),
              result,
              constraint,
              hierarchiesInRelation);
        }
      } else {
        for (AnnotationMirror top : tops) {
          final AnnotationMirror tAnno = typeT.getAnnotationInHierarchy(top);

          if (tAnno == null) {
            hierarchiesInRelation.add(top);
          }
        }

        addToTypeRelationship(
            (TypeVariable) TypeAnnotationUtils.unannotatedType(typeT.getUnderlyingType()),
            typeU,
            result,
            constraint,
            hierarchiesInRelation);
      }
    }

    return result;
  }

  private void addToTargetRelationship(
      TypeVariable typeT,
      TypeVariable typeU,
      ConstraintMap result,
      TUConstraint constraint,
      AnnotationMirrorSet hierarchiesInRelation) {
    if (constraint instanceof TIsU) {
      result.addTargetEquality(typeT, typeU, hierarchiesInRelation);
    } else if (constraint instanceof TSuperU) {
      result.addTargetSupertype(typeT, typeU, hierarchiesInRelation);
    } else {
      result.addTargetSubtype(typeT, typeU, hierarchiesInRelation);
    }
  }

  public void addToPrimaryRelationship(
      TypeVariable typeVariable,
      TUConstraint constraint,
      ConstraintMap result,
      AnnotationMirrorSet annotationMirrors,
      QualifierHierarchy qualifierHierarchy) {
    if (constraint instanceof TIsU) {
      result.addPrimaryEqualities(typeVariable, qualifierHierarchy, annotationMirrors);
    } else if (constraint instanceof TSuperU) {
      result.addPrimarySupertype(typeVariable, qualifierHierarchy, annotationMirrors);
    } else {
      result.addPrimarySubtypes(typeVariable, qualifierHierarchy, annotationMirrors);
    }
  }

  public void addToTypeRelationship(
      TypeVariable target,
      AnnotatedTypeMirror type,
      ConstraintMap result,
      TUConstraint constraint,
      AnnotationMirrorSet hierarchies) {
    if (constraint instanceof TIsU) {
      result.addTypeEqualities(target, type, hierarchies);
    } else if (constraint instanceof TSuperU) {
      result.addTypeSupertype(target, type, hierarchies);
    } else {
      result.addTypeSubtype(target, type, hierarchies);
    }
  }
}