package com.twitter.elephantbird.crunch;

import com.google.protobuf.Message;
import com.twitter.elephantbird.mapreduce.io.BinaryWritable;
import com.twitter.elephantbird.mapreduce.io.ProtobufConverter;
import com.twitter.elephantbird.mapreduce.io.ProtobufWritable;
import com.twitter.elephantbird.mapreduce.io.ThriftConverter;
import com.twitter.elephantbird.mapreduce.io.ThriftWritable;
import org.apache.crunch.MapFn;
import org.apache.crunch.types.PType;
import org.apache.crunch.types.writable.Writables;
import org.apache.thrift.TBase;

/**
 * Crunch {@link org.apache.crunch.types.PType} implementations for Elephant Bird's {@code Writable} types.
 */
public class EBTypes {

  public static <M extends Message> PType<M> protos(Class<M> protoClass) {
    return Writables.derived(protoClass, new ProtoInFn(protoClass), new ProtoOutFn(protoClass),
        Writables.writables(ProtobufWritable.class));
  }

  public static <M extends TBase<?, ?>> PType<M> thrifts(Class<M> thriftClass) {
    return Writables.derived(thriftClass, new ThriftInFn(thriftClass), new ThriftOutFn(thriftClass),
        Writables.writables(ThriftWritable.class));
  }

  private static class ProtoInFn<T extends Message> extends MapFn<ProtobufWritable<T>, T> {
    private final Class<T> clazz;
    private transient ProtobufConverter<T> converter;

    public ProtoInFn(Class<T> clazz) {
      this.clazz = clazz;
    }

    @Override
    public void initialize() {
      this.converter = ProtobufConverter.newInstance(clazz);
    }

    @Override
    public T map(ProtobufWritable<T> input) {
      input.setConverter(converter);
      return input.get();
    }
  }

  private static class ProtoOutFn<T extends Message> extends MapFn<T, ProtobufWritable<T>> {
    private final Class<T> clazz;

    public ProtoOutFn(Class<T> clazz) {
      this.clazz = clazz;
    }

    @Override
    public ProtobufWritable<T> map(T input) {
      ProtobufWritable<T> w = ProtobufWritable.newInstance(clazz);
      w.set(input);
      return w;
    }
  }

  private static class ThriftInFn<T extends TBase<?, ?>> extends MapFn<ThriftWritable<T>, T> {
    private final Class<T> clazz;
    private transient ThriftConverter<T> converter;

    public ThriftInFn(Class<T> clazz) {
      this.clazz = clazz;
    }

    @Override
    public void initialize() {
      this.converter = ThriftConverter.newInstance(clazz);
    }

    @Override
    public T map(ThriftWritable<T> input) {
      input.setConverter(converter);
      return input.get();
    }
  }

  private static class ThriftOutFn<T extends TBase<?, ?>> extends MapFn<T, ThriftWritable<T>> {
    private final Class<T> clazz;

    public ThriftOutFn(Class<T> clazz) {
      this.clazz = clazz;
    }

    @Override
    public ThriftWritable<T> map(T input) {
      ThriftWritable<T> w = ThriftWritable.newInstance(clazz);
      w.set(input);
      return w;
    }
  }
}

@@ -1763,32 +1763,86 @@
  Class%3CT%3E clazz;
+%0A    private transient ProtobufConverter%3CT%3E converter;
 %0A%0A    public Pro
@@ -1892,32 +1892,145 @@
 = clazz;%0A    %7D%0A%0A
+    @Override%0A    public void initialize() %7B%0A      this.converter = ProtobufConverter.newInstance(clazz);%0A    %7D%0A%0A
     @Override%0A  
@@ -2102,16 +2102,20 @@
 e%3CT%3E w =
+ new
  Protobu
@@ -2123,34 +2123,53 @@
 Writable
-.newInstance(clazz
+%3CT%3E();%0A      w.setConverter(converter
 );%0A     
@@ -2837,16 +2837,68 @@
 %3E clazz;
+%0A    private transient ThriftConverter%3CT%3E converter;
 %0A%0A    pu
@@ -2957,32 +2957,143 @@
 = clazz;%0A    %7D%0A%0A
+    @Override%0A    public void initialize() %7B%0A      this.converter = ThriftConverter.newInstance(clazz);%0A    %7D%0A%0A
     @Override%0A  
@@ -3161,16 +3161,20 @@
 e%3CT%3E w =
+ new
  ThriftW
@@ -3180,34 +3180,53 @@
 Writable
-.newInstance(clazz
+%3CT%3E();%0A      w.setConverter(converter
 );%0A     

Use converter instance directly for ProtoOutFn and ThriftOutFn