/*
 * Copyright (c) 2002,2006 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.shrike.shrikeCT;

import com.ibm.wala.util.collections.HashMapFactory;
import com.ibm.wala.util.collections.Pair;
import com.ibm.wala.util.debug.Assertions;
import java.util.Arrays;
import java.util.Map;

/**
 * This class reads Annotations attributes, e.g., RuntimeInvisibleAnnotations.
 *
 * @author sjfink
 */
public class AnnotationsReader extends AttributeReader {

  /** offset in class file where this attribute begins */
  protected final int beginOffset;

  public AnnotationsReader(ClassReader.AttrIterator iter, String label)
      throws InvalidClassFileException {
    super(iter, label);
    beginOffset = attr;
  }

  /** @return number of annotations in this attribute */
  public int getAnnotationCount() throws InvalidClassFileException {
    int offset = beginOffset + 6;
    checkSize(offset, 2);
    return cr.getUShort(offset);
  }

  /** @return total length of this attribute in bytes, <b>including</b> the first 6 bytes */
  public int getAttributeSize() throws InvalidClassFileException {
    int offset = beginOffset + 2;
    checkSize(offset, 4);
    return cr.getInt(offset) + 6;
  }

  /**
   * get the Utf8 constant pool value, where the constant pool offset is given in the class
   *
   * @param offset offset in the class file at which the constant pool offset is given
   */
  protected String getUtf8ConstantPoolValue(int offset) throws InvalidClassFileException {
    checkSize(offset, 2);
    int cpOffset = cr.getUShort(offset);
    return cr.getCP().getCPUtf8(cpOffset);
  }

  /**
   * Marker interface for possible element values in an annotation attribute.
   *
   * @see AnnotationsReader#readElementValueAndSize(int)
   */
  public interface ElementValue {}

  /**
   * Represents a constant argument to an annotation. Class arguments (e.g., {@code Foo.class}) are
   * also represented with this type, with the value being the String class name.
   */
  public static class ConstantElementValue implements ElementValue {

    /** the constant value */
    public final Object val;

    public ConstantElementValue(Object val) {
      this.val = val;
    }

    @Override
    public String toString() {
      return String.valueOf(val);
    }

    @Override
    public int hashCode() {
      final int prime = 31;
      int result = 1;
      result = prime * result + ((val == null) ? 0 : val.hashCode());
      return result;
    }

    @Override
    public boolean equals(Object obj) {
      if (this == obj) return true;
      if (obj == null) return false;
      if (getClass() != obj.getClass()) return false;
      ConstantElementValue other = (ConstantElementValue) obj;
      if (val == null) {
        if (other.val != null) return false;
      } else if (!val.equals(other.val)) return false;
      return true;
    }
  }

  /** Represents enum constant annotation arguments. */
  public static class EnumElementValue implements ElementValue {

    /** the name of the enum type */
    public final String enumType;

    /** the enum value */
    public final String enumVal;

    public EnumElementValue(String enumType, String enumVal) {
      super();
      this.enumType = enumType;
      this.enumVal = enumVal;
    }

    @Override
    public String toString() {
      return "EnumElementValue [type=" + enumType + ", val=" + enumVal + ']';
    }
  }

  /** represents an annotation argument that itself is an array of arguments */
  public static class ArrayElementValue implements ElementValue {

    /** the values contained in the array */
    public final ElementValue[] vals;

    public ArrayElementValue(ElementValue[] vals) {
      super();
      this.vals = vals;
    }

    @Override
    public String toString() {
      return "ArrayElementValue [vals=" + Arrays.toString(vals) + ']';
    }
  }

  /**
   * get all the annotations declared in this attribute.
   *
   * <pre>
   * RuntimeVisibleAnnotations_attribute {
   *   u2         attribute_name_index;
   *   u4         attribute_length;
   *   u2         num_annotations;
   *   annotation annotations[num_annotations];
   * }
   * </pre>
   *
   * @see <a href="https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-4.html#jvms-4.7.16">JLS
   *     (SE8), 4.7.16</a>
   */
  public AnnotationAttribute[] getAllAnnotations() throws InvalidClassFileException {
    AnnotationAttribute[] result = new AnnotationAttribute[getAnnotationCount()];
    int offset = beginOffset + 8; // skip attribute_name_index,
    // attribute_length, and num_annotations
    for (int i = 0; i < result.length; i++) {
      Pair<AnnotationAttribute, Integer> attributeAndSize = getAttributeAndSize(offset);
      result[i] = attributeAndSize.fst;
      offset += attributeAndSize.snd;
    }
    return result;
  }

  /**
   *
   *
   * <pre>
   * param_annotations {
   *   u2 attribute_name_index;
   *   u4 attribute_length;
   *   u1 num_parameters;
   *   {
   *     u2 num_annotations;
   *     annotation annotations[num_annotations];
   *   } parameter_annotations[num_parameters];
   * </pre>
   */
  public AnnotationAttribute[][] getAllParameterAnnotations() throws InvalidClassFileException {
    int numParamOffset = beginOffset + 6;
    checkSize(numParamOffset, 1);
    int paramCount = cr.getByte(numParamOffset);
    AnnotationAttribute[][] result = new AnnotationAttribute[paramCount][];
    // skip attribute_name_index, attribute_length, and num_parameters
    int offset = beginOffset + 7;
    for (int i = 0; i < result.length; i++) {
      checkSize(offset, 2);
      result[i] = new AnnotationAttribute[cr.getUShort(offset)];
      offset += 2;
      for (int j = 0; j < result[i].length; j++) {
        Pair<AnnotationAttribute, Integer> attributeAndSize = getAttributeAndSize(offset);
        result[i][j] = attributeAndSize.fst;
        offset += attributeAndSize.snd;
      }
    }
    return result;
  }

  /**
   *
   *
   * <pre>
   * annotation {
   *   u2 type_index;
   *   u2 num_element_value_pairs;
   *   { u2 element_name_index;
   *     element_value value;
   *   } element_value_pairs[num_element_value_pairs]
   * </pre>
   */
  protected Pair<AnnotationAttribute, Integer> getAttributeAndSize(int begin)
      throws InvalidClassFileException {
    String type = getUtf8ConstantPoolValue(begin);
    int numElementValuePairs = cr.getUShort(begin + 2);
    int size = 4;
    int offset = begin + 4;
    Map<String, ElementValue> elementName2Val = HashMapFactory.make();
    for (int i = 0; i < numElementValuePairs; i++) {
      String elementName = getUtf8ConstantPoolValue(offset);
      offset += 2;
      Pair<ElementValue, Integer> elementValAndSize = readElementValueAndSize(offset);
      offset += elementValAndSize.snd;
      size += elementValAndSize.snd + 2;
      elementName2Val.put(elementName, elementValAndSize.fst);
    }
    return Pair.make(new AnnotationAttribute(type, elementName2Val), size);
  }

  /**
   * Representation of an annotation attribute. An annotation has the following format in the
   * bytecode:
   *
   * <pre>
   * annotation {
   *   u2 type_index;
   *   u2 num_element_value_pairs;
   *   {  u2 element_name_index;
   *      element_value value;
   * } element_value_pairs[num_element_value_pairs];
   * </pre>
   *
   * See the JVM specification section 4.7.16 for details.
   *
   * <p>This class implements {@link ElementValue} to handle nested annotations.
   */
  public static class AnnotationAttribute implements ElementValue {

    /** the type of the annotation */
    public final String type;

    /** the arguments to the annotation */
    public final Map<String, ElementValue> elementValues;

    public AnnotationAttribute(String type, Map<String, ElementValue> elementValues) {
      super();
      this.type = type;
      this.elementValues = elementValues;
    }

    @Override
    public String toString() {
      return "AnnotationElementValue [type=" + type + ", elementValues=" + elementValues + ']';
    }
  }

  /**
   *
   *
   * <pre>
   * element_value {
   *   u1 tag;
   *   union {
   *     u2 const_value_index;
   *     {  u2 type_name_index;
   *        u2 const_name_index;
   *     } enum_const_value;
   *     u2 class_info_index;
   *     annotation annotation_value;
   *     {  u2 num_values;
   *        element_value values[num_values];
   *     } array_value;
   *   } value;
   * </pre>
   *
   * A constant value (including class info) is represented by a {@link ConstantElementValue}. An
   * enum constant value is represented by an {@link EnumElementValue}. An array value is
   * represented by an {@link ArrayElementValue}. Finally, a nested annotation is represented by an
   * {@link AnnotationAttribute}.
   */
  protected Pair<ElementValue, Integer> readElementValueAndSize(int offset)
      throws IllegalArgumentException, InvalidClassFileException {
    char tag = (char) cr.getByte(offset);
    // meaning of this short depends on the tag
    int nextShort = cr.getUShort(offset + 1);
    switch (tag) {
      case 'B':
      case 'C':
      case 'I':
      case 'S':
      case 'Z':
        return Pair.<ElementValue, Integer>make(
            new ConstantElementValue(cr.getCP().getCPInt(nextShort)), 3);
      case 'J':
        return Pair.<ElementValue, Integer>make(
            new ConstantElementValue(cr.getCP().getCPLong(nextShort)), 3);
      case 'D':
        return Pair.<ElementValue, Integer>make(
            new ConstantElementValue(cr.getCP().getCPDouble(nextShort)), 3);
      case 'F':
        return Pair.<ElementValue, Integer>make(
            new ConstantElementValue(cr.getCP().getCPFloat(nextShort)), 3);
      case 's': // string
      case 'c': // class; just represent as a constant element with the type name
        return Pair.<ElementValue, Integer>make(
            new ConstantElementValue(cr.getCP().getCPUtf8(nextShort)), 3);
      case 'e': // enum
        return Pair.<ElementValue, Integer>make(
            new EnumElementValue(
                cr.getCP().getCPUtf8(nextShort), cr.getCP().getCPUtf8(cr.getUShort(offset + 3))),
            5);
      case '[': // array
        int numValues = nextShort;
        int numArrayBytes = 3; // start with 3 for the tag and num_values bytes
        ElementValue[] vals = new ElementValue[numValues];
        // start curOffset at beginning of array values
        int curArrayOffset = offset + 3;
        for (int i = 0; i < numValues; i++) {
          Pair<ElementValue, Integer> arrayElemValueAndSize =
              readElementValueAndSize(curArrayOffset);
          vals[i] = arrayElemValueAndSize.fst;
          curArrayOffset += arrayElemValueAndSize.snd;
          numArrayBytes += arrayElemValueAndSize.snd;
        }
        return Pair.<ElementValue, Integer>make(new ArrayElementValue(vals), numArrayBytes);
      case '@': // annotation
        Pair<AnnotationAttribute, Integer> attributeAndSize = getAttributeAndSize(offset + 1);
        // add 1 to size for the tag
        return Pair.<ElementValue, Integer>make(attributeAndSize.fst, attributeAndSize.snd + 1);
      default:
        assert false;
        return null;
    }
  }

  // //////////////
  // utility methods for reading well-known annotation types
  // //////////////

  public enum AnnotationType {
    RuntimeVisibleAnnotations,
    RuntimeInvisibleAnnotations,
    RuntimeVisibleParameterAnnotations,
    RuntimeInvisibleParameterAnnotations
  }

  public static boolean isKnownAnnotation(String name) {
    for (AnnotationType t : AnnotationType.values()) {
      if (t.name().equals(name)) {
        return true;
      }
    }
    return false;
  }

  public static AnnotationsReader getReaderForAnnotation(
      AnnotationType type, ClassReader.AttrIterator iter) {
    // search for the desired attribute
    final String attrName = type.toString();
    try {
      for (; iter.isValid(); iter.advance()) {
        if (iter.getName().equals(attrName)) {
          return new AnnotationsReader(iter, attrName);
        }
      }
    } catch (InvalidClassFileException e) {
      Assertions.UNREACHABLE();
    }
    return null;
  }
}
