package org.checkerframework.framework.util.element;

import com.sun.tools.javac.code.Attribute;
import com.sun.tools.javac.code.Attribute.TypeCompound;
import com.sun.tools.javac.code.Symbol;
import com.sun.tools.javac.code.TargetType;
import com.sun.tools.javac.code.TypeAnnotationPosition;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import javax.lang.model.element.Element;
import org.checkerframework.framework.type.AnnotatedTypeFactory;
import org.checkerframework.framework.type.AnnotatedTypeMirror;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedExecutableType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedTypeVariable;
import org.checkerframework.framework.util.element.ElementAnnotationUtil.UnexpectedAnnotationLocationException;
import org.checkerframework.javacutil.BugInCF;
import org.checkerframework.javacutil.ElementUtils;
import org.plumelib.util.StringsPlume;

/**
 * Adds annotations from element to the return type, formal parameter types, type parameters, and
 * throws clauses of the AnnotatedExecutableType type.
 */
public class MethodApplier extends TargetedElementAnnotationApplier {

  /** Apply annotations from {@code element} to {@code type}. */
  public static void apply(
      AnnotatedTypeMirror type, Element element, AnnotatedTypeFactory typeFactory)
      throws UnexpectedAnnotationLocationException {
    new MethodApplier(type, element, typeFactory).extractAndApply();
  }

  public static boolean accepts(final AnnotatedTypeMirror typeMirror, final Element element) {
    return element instanceof Symbol.MethodSymbol && typeMirror instanceof AnnotatedExecutableType;
  }

  private final AnnotatedTypeFactory typeFactory;

  /** Method being annotated, this symbol contains all relevant annotations. */
  private final Symbol.MethodSymbol methodSymbol;

  private final AnnotatedExecutableType methodType;

  MethodApplier(AnnotatedTypeMirror type, Element element, AnnotatedTypeFactory typeFactory) {
    super(type, element);
    this.typeFactory = typeFactory;
    this.methodSymbol = (Symbol.MethodSymbol) element;
    this.methodType = (AnnotatedExecutableType) type;
  }

  /**
   * Returns receiver, returns, and throws. See extract and apply as we also annotate type params.
   *
   * @return receiver, returns, and throws
   */
  @Override
  protected TargetType[] annotatedTargets() {
    return new TargetType[] {
      TargetType.METHOD_RECEIVER, TargetType.METHOD_RETURN, TargetType.THROWS
    };
  }

  /**
   * Returns all possible annotation positions for a method except those in annotatedTargets.
   *
   * @return all possible annotation positions for a method except those in annotatedTargets
   */
  @Override
  protected TargetType[] validTargets() {
    return new TargetType[] {
      TargetType.LOCAL_VARIABLE,
      TargetType.RESOURCE_VARIABLE,
      TargetType.EXCEPTION_PARAMETER,
      TargetType.NEW,
      TargetType.CAST,
      TargetType.INSTANCEOF,
      TargetType.METHOD_INVOCATION_TYPE_ARGUMENT,
      TargetType.CONSTRUCTOR_INVOCATION_TYPE_ARGUMENT,
      TargetType.METHOD_REFERENCE,
      TargetType.CONSTRUCTOR_REFERENCE,
      TargetType.METHOD_REFERENCE_TYPE_ARGUMENT,
      TargetType.CONSTRUCTOR_REFERENCE_TYPE_ARGUMENT,
      TargetType.METHOD_TYPE_PARAMETER,
      TargetType.METHOD_TYPE_PARAMETER_BOUND,
      TargetType.METHOD_FORMAL_PARAMETER,
      // TODO: from generic anonymous classes; remove when
      // we can depend on only seeing classfiles that were
      // generated by a javac that contains a fix for:
      // https://bugs.openjdk.java.net/browse/JDK-8198945
      TargetType.CLASS_EXTENDS,
      // TODO: Test case from Issue 3277 produces invalid position.
      // Ignore until this javac bug is fixed:
      // https://bugs.openjdk.java.net/browse/JDK-8233945
      TargetType.UNKNOWN
    };
  }

  /**
   * Returns the annotations on the method symbol (element).
   *
   * @return the annotations on the method symbol (element)
   */
  @Override
  protected Iterable<Attribute.TypeCompound> getRawTypeAttributes() {
    return methodSymbol.getRawTypeAttributes();
  }

  @Override
  protected boolean isAccepted() {
    return MethodApplier.accepts(type, element);
  }

  /**
   * Sets the method's element, annotates its return type, parameters, type parameters, and throws
   * annotations.
   */
  @Override
  public void extractAndApply() throws UnexpectedAnnotationLocationException {
    methodType.setElement(methodSymbol); // Preserves previous behavior

    // Add declaration annotations to the return type if
    if (methodType.getReturnType() instanceof AnnotatedTypeVariable) {
      applyTypeVarUseOnReturnType();
    }
    ElementAnnotationUtil.addDeclarationAnnotationsFromElement(
        methodType.getReturnType(), methodSymbol.getAnnotationMirrors());

    final List<AnnotatedTypeMirror> params = methodType.getParameterTypes();
    for (int i = 0; i < params.size(); ++i) {
      // Add declaration annotations to the parameter type
      ElementAnnotationUtil.addDeclarationAnnotationsFromElement(
          params.get(i), methodSymbol.getParameters().get(i).getAnnotationMirrors());
    }

    // ensures that we check that there are only valid target types on this class, there are no
    // "invalid" locations
    super.extractAndApply();

    ElementAnnotationUtil.applyAllElementAnnotations(
        methodType.getParameterTypes(), methodSymbol.getParameters(), typeFactory);
    ElementAnnotationUtil.applyAllElementAnnotations(
        methodType.getTypeVariables(), methodSymbol.getTypeParameters(), typeFactory);
  }

  // NOTE that these are the only locations not handled elsewhere, otherwise we call apply
  @Override
  protected void handleTargeted(final List<TypeCompound> targeted)
      throws UnexpectedAnnotationLocationException {
    final List<TypeCompound> unmatched = new ArrayList<>();
    final Map<TargetType, List<TypeCompound>> targetTypeToAnno =
        ElementAnnotationUtil.partitionByTargetType(
            targeted,
            unmatched,
            TargetType.METHOD_RECEIVER,
            TargetType.METHOD_RETURN,
            TargetType.THROWS);

    ElementAnnotationUtil.annotateViaTypeAnnoPosition(
        methodType.getReceiverType(), targetTypeToAnno.get(TargetType.METHOD_RECEIVER));
    ElementAnnotationUtil.annotateViaTypeAnnoPosition(
        methodType.getReturnType(), targetTypeToAnno.get(TargetType.METHOD_RETURN));
    applyThrowsAnnotations(targetTypeToAnno.get(TargetType.THROWS));

    if (!unmatched.isEmpty()) {
      throw new BugInCF(
          "Unexpected annotations ( "
              + StringsPlume.join(",", unmatched)
              + " ) for"
              + "type ( "
              + type
              + " ) and element ( "
              + element
              + " ) ");
    }
  }

  /** For each thrown type, collect all the annotations for that type and apply them. */
  private void applyThrowsAnnotations(final List<Attribute.TypeCompound> annos)
      throws UnexpectedAnnotationLocationException {
    final List<AnnotatedTypeMirror> thrown = methodType.getThrownTypes();
    if (thrown.isEmpty()) {
      return;
    }

    Map<AnnotatedTypeMirror, List<TypeCompound>> typeToAnnos = new LinkedHashMap<>();
    for (final AnnotatedTypeMirror thrownType : thrown) {
      typeToAnnos.put(thrownType, new ArrayList<>());
    }

    for (TypeCompound anno : annos) {
      final TypeAnnotationPosition annoPos = anno.position;
      if (annoPos.type_index >= 0 && annoPos.type_index < thrown.size()) {
        final AnnotatedTypeMirror thrownType = thrown.get(annoPos.type_index);
        typeToAnnos.get(thrownType).add(anno);

      } else {
        throw new BugInCF(
            "MethodApplier.applyThrowsAnnotation: "
                + "invalid throws index "
                + annoPos.type_index
                + " for annotation: "
                + anno
                + " for element: "
                + ElementUtils.getQualifiedName(element));
      }
    }

    for (final Map.Entry<AnnotatedTypeMirror, List<TypeCompound>> typeToAnno :
        typeToAnnos.entrySet()) {
      ElementAnnotationUtil.annotateViaTypeAnnoPosition(typeToAnno.getKey(), typeToAnno.getValue());
    }
  }

  /**
   * If the return type is a use of a type variable first apply the bound annotations from the type
   * variables declaration.
   */
  private void applyTypeVarUseOnReturnType() throws UnexpectedAnnotationLocationException {
    new TypeVarUseApplier(methodType.getReturnType(), methodSymbol, typeFactory).extractAndApply();
  }
}