package com.google.cloud.dataflow.contrib.hadoop;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.cloud.dataflow.sdk.coders.AvroCoder;
import com.google.cloud.dataflow.sdk.coders.Coder;
import com.google.cloud.dataflow.sdk.coders.CoderException;
import com.google.cloud.dataflow.sdk.coders.StandardCoder;
import com.google.cloud.dataflow.sdk.util.CloudObject;
import com.google.cloud.dataflow.sdk.util.PropertyNames;
import com.google.common.base.Preconditions;
import org.apache.avro.mapred.AvroWrapper;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Collections;
import java.util.List;

public class AvroWrapperCoder<W extends AvroWrapper<D>, D> extends StandardCoder<W> {
  private static final long serialVersionUID = 0L;

  private final Class<W> wrapperType;
  private final AvroCoder<D> datumCoder;

  private AvroWrapperCoder(Class<W> wrapperType, AvroCoder<D> datumCoder) {
    this.wrapperType = wrapperType;
    this.datumCoder = datumCoder;
  }

  public static <W extends AvroWrapper<D>, D>
  AvroWrapperCoder<W, D>of(Class<W> wrapperType, AvroCoder<D> datumCoder) {
    return new AvroWrapperCoder<>(wrapperType, datumCoder);
  }

  @JsonCreator
  @SuppressWarnings("unchecked")
  public static AvroWrapperCoder<?, ?> of(
      @JsonProperty("wrapperType") String wrapperType,
      @JsonProperty(PropertyNames.COMPONENT_ENCODINGS) List<Coder<?>> components)
      throws ClassNotFoundException {
    Class<?> clazz = Class.forName(wrapperType);
    if (!AvroWrapper.class.isAssignableFrom(clazz)) {
      throw new ClassNotFoundException(
          "Class " + wrapperType + " does not implement AvroWrapper");
    }
    Preconditions.checkArgument(components.size() == 1,
        "Expecting 1 component, got " + components.size());
    return of((Class<? extends AvroWrapper>) clazz, (AvroCoder<?>) components.get(0));
  }

  @Override
  public void encode(W value, OutputStream outStream, Context context) throws IOException {
    datumCoder.encode(value.datum(), outStream, context);
  }

  @Override
  public W decode(InputStream inStream, Context context) throws IOException {
    try {
      W wrapper = wrapperType.newInstance();
      wrapper.datum(datumCoder.decode(inStream, context));
      return wrapper;
    } catch (InstantiationException | IllegalAccessException e) {
      throw new CoderException("unable to deserialize record", e);
    }
  }

  @Override
  public List<? extends Coder<?>> getCoderArguments() {
    return Collections.singletonList(datumCoder);
  }

  @Override
  public CloudObject asCloudObject() {
    CloudObject result = super.asCloudObject();
    result.put("wrapperType", wrapperType.getName());
    return result;
  }

  @Override
  public void verifyDeterministic() throws NonDeterministicException {}

}

@@ -716,16 +716,267 @@
 .List;%0A%0A
+/**%0A * A %7B@code AvroWrapperCoder%7D is a %7B@link com.google.cloud.dataflow.sdk.coders.Coder%7D for a Java%0A * class that implements %7B@link org.apache.avro.mapred.AvroWrapper%7D.%0A *%0A * @param %3CW%3E the type of the wrapper%0A * @param %3CD%3E the type of the datum%0A */%0A
 public c
@@ -1336,16 +1336,186 @@
 r;%0A  %7D%0A%0A
+  /**%0A   * Return a %7B@code AvroWrapperCoder%7D instance for the provided element class.%0A   * @param %3CW%3E the type of the wrapper%0A   * @param %3CD%3E the type of the datum%0A   */%0A
   public

add JavaDoc to AvroWrapperCoder