package org.checkerframework.checker.units;

import java.util.Collection;
import java.util.Collections;
import java.util.Set;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.util.Elements;
import org.checkerframework.checker.units.qual.Prefix;
import org.checkerframework.common.basetype.BaseTypeChecker;
import org.checkerframework.framework.type.DefaultAnnotatedTypeFormatter;
import org.checkerframework.framework.util.AnnotationFormatter;
import org.checkerframework.framework.util.DefaultAnnotationFormatter;
import org.checkerframework.javacutil.AnnotationUtils;

public class UnitsAnnotatedTypeFormatter extends DefaultAnnotatedTypeFormatter {
  protected final BaseTypeChecker checker;
  protected final Elements elements;

  public UnitsAnnotatedTypeFormatter(BaseTypeChecker checker) {
    // Utilize the Default Type Formatter, but force it to print out Invisible Qualifiers
    // keep super call in sync with implementation in DefaultAnnotatedTypeFormatter
    // keep checker options in sync with implementation in AnnotatedTypeFactory
    super(
        new UnitsFormattingVisitor(
            checker,
            new UnitsAnnotationFormatter(checker),
            checker.hasOption("printVerboseGenerics"),
            true));

    this.checker = checker;
    this.elements = checker.getElementUtils();
  }

  protected static class UnitsFormattingVisitor
      extends DefaultAnnotatedTypeFormatter.FormattingVisitor {
    protected final BaseTypeChecker checker;
    protected final Elements elements;

    public UnitsFormattingVisitor(
        BaseTypeChecker checker,
        AnnotationFormatter annoFormatter,
        boolean printVerboseGenerics,
        boolean defaultInvisiblesSetting) {

      super(annoFormatter, printVerboseGenerics, defaultInvisiblesSetting);
      this.checker = checker;
      this.elements = checker.getElementUtils();
    }
  }

  /** Format the error printout of any units qualifier that uses Prefix.one. */
  protected static class UnitsAnnotationFormatter extends DefaultAnnotationFormatter {
    protected final BaseTypeChecker checker;
    protected final Elements elements;

    public UnitsAnnotationFormatter(BaseTypeChecker checker) {
      this.checker = checker;
      this.elements = checker.getElementUtils();
    }

    @Override
    public String formatAnnotationString(
        Collection<? extends AnnotationMirror> annos, boolean printInvisible) {
      // create an empty annotation set
      Set<AnnotationMirror> trimmedAnnoSet = AnnotationUtils.createAnnotationSet();

      // loop through all the annotation mirrors to see if they use Prefix.one, remove
      // Prefix.one if it does
      for (AnnotationMirror anno : annos) {
        if (UnitsRelationsTools.getPrefix(anno) == Prefix.one) {
          anno = UnitsRelationsTools.removePrefix(elements, anno);
        }
        // add to set
        trimmedAnnoSet.add(anno);
      }

      return super.formatAnnotationString(
          Collections.unmodifiableSet(trimmedAnnoSet), printInvisible);
    }
  }
}