package org.apache.hadoop.io;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.EnumSet;

/**
 * A Writable implementation that can store enumsets.
 */
public class EnumSetWritable<E extends Enum<E>> implements Writable {

  private int size;

  private long data = 0;

  public EnumSetWritable() {
    size = 1;
  }

  public EnumSetWritable(Class<Enum<E>> type) {

    int size = type.getEnumConstants().length;
    if (size > 30) {
      throw new IllegalArgumentException(
          "Maximum number of bytes is 4 (30 bits)");
    }
    this.size = size / 8;
    if (size == 0)
      size = 1;
  }

  public EnumSetWritable(EnumSet<E> value) {
    putEnumSet(value);
  }

  /**
   * Gets enumset.
   * @param <E>
   * @param enumClass
   * @return
   */
  public EnumSet<E> getEnumSet(Class<E> enumClass) {
    return toEnumSet(enumClass, data);
  }

  private EnumSet<E> toEnumSet(Class<E> enumClass, final long vector) {
    EnumSet<E> set = EnumSet.noneOf(enumClass);
    long mask = 1;
    for (E e : enumClass.getEnumConstants()) {
      if ((mask & vector) == mask) {
        set.add(e);
      }
      mask <<= 1;
    }
    return set;
  }

  public void putEnumSet(EnumSet<E> set) {
    data = toLong(set);
  }

  private long toLong(EnumSet<E> s) {
    long vector = 0;
    int max = 0;

    for (E e : s) {
      vector += (1 << e.ordinal());
      if (max < e.ordinal()) {
        max = e.ordinal();
      }
    }
    size = 1 + (max + 2) / 8;
    return vector;
  }

  public void readFields(DataInput in) throws IOException {

    reset();
    int bytes = 0;
    int roll = 0;

    do {
      int read = in.readByte() & 0xff;

      if (size == 0) {
        size = (read & 0x03);
      }

      data += ((long) read << roll);

      roll += 8;
    } while (++bytes < size);

    //remove size bits
    data >>= 2;
  }

  /**
   * Reset state.
   */
  private void reset() {
    data = 0;
    size = 0;
  }

  public void write(DataOutput out) throws IOException {
    long shifted = (size & 0x03) + (this.data << 2);
    for (int i = 0; i < size; i++) {
      int data = (int) (shifted & 0xff);
      shifted >>= 8;
      out.writeByte(data);
    }
  }

  @Override
  public boolean equals(Object obj) {
    if (!(obj instanceof EnumSetWritable)) {
      return super.equals(obj);
    }
    EnumSetWritable that = (EnumSetWritable) obj;
    return this.data == that.data;
  }

}