package com.twitter.elephantbird.pig.load;

import java.io.IOException;

import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.pig.ResourceSchema;
import org.apache.pig.data.Tuple;
import org.apache.thrift.TBase;

import com.google.protobuf.Message;
import com.google.protobuf.Descriptors.Descriptor;
import com.twitter.elephantbird.mapreduce.input.LzoRecordReader;
import com.twitter.elephantbird.mapreduce.input.MultiInputFormat;
import com.twitter.elephantbird.mapreduce.io.BinaryWritable;
import com.twitter.elephantbird.pig.util.PigUtil;
import com.twitter.elephantbird.pig.util.ProtobufToPig;
import com.twitter.elephantbird.pig.util.ProtobufTuple;
import com.twitter.elephantbird.pig.util.ThriftToPig;
import com.twitter.elephantbird.util.Protobufs;
import com.twitter.elephantbird.util.TypeRef;

/**
 * A loader based on {@link MultiInputFormat} to load input with
 * different file formats.
 *
 * @see MultiInputFormat
 */
public class MultiFormatLoader<M> extends LzoBaseLoadFunc {

  private TypeRef<M> typeRef = null;
  private ObjToTuple tupleImpl;

  // Pig requires a default constructor. Do not use.
  public MultiFormatLoader() {
  }

  /**
   * @param className Thrift or Protobuf class
   */
  public MultiFormatLoader(String className) {
    Class<?> clazz = PigUtil.getClass(className);
    typeRef = new TypeRef<M>(clazz){};

    // initialize tupleImpl

    if (Message.class.isAssignableFrom(clazz)) {
      tupleImpl = new ProtobufToTuple(typeRef);
    } else if (TBase.class.isAssignableFrom(clazz)) {
      tupleImpl = new ThriftToTuple(typeRef);
    } else {
      throw new RuntimeException(className + " is not a Protobuf or Thrift class");
    }
  }

  /**
   * Return next Tuple from input.
   * <p>
   * A small fraction of bad records in input are tolerated.
   * See  {@link LzoRecordReader} for more information on error handling.
   */
  public Tuple getNext() throws IOException {
    M value = getNextBinaryValue(typeRef);

    return value != null ?
        tupleImpl.toTuple(value) : null;
  }

  @Override
  public ResourceSchema getSchema(String location, Job job) throws IOException {
    return tupleImpl.getSchema();
  }

  @Override
  public InputFormat<LongWritable, BinaryWritable<M>> getInputFormat() throws IOException {
    return new MultiInputFormat<M>(typeRef);
  }

  // END OF Pig Loader implementation

  /*
   * classes to convert to Thrift or Protobuf object into tuples:
   */

  private static interface ObjToTuple {
    abstract ResourceSchema getSchema();
    abstract Tuple toTuple(Object obj);
  };

  private static class ThriftToTuple implements ObjToTuple {
    private TypeRef<TBase<?, ?>> typeRef;
    private ThriftToPig<TBase<?, ?>> thriftToPig;

    @SuppressWarnings("unchecked")
    ThriftToTuple(TypeRef<?> typeRef) {
      this.typeRef = (TypeRef<TBase<?, ?>>) typeRef;
      this.thriftToPig = ThriftToPig.newInstance(this.typeRef);
    }

    @Override
    public ResourceSchema getSchema() {
      return new ResourceSchema(thriftToPig.toSchema());
    }

    @Override
    public Tuple toTuple(Object obj) {
      return thriftToPig.getLazyTuple((TBase<?, ?>)obj);
    }
  }

  private static class ProtobufToTuple implements ObjToTuple {
    private TypeRef<Message> typeRef;

    @SuppressWarnings("unchecked")
    ProtobufToTuple(TypeRef<?> typeRef) {
      this.typeRef = (TypeRef<Message>)typeRef;
    }

    @Override
    public ResourceSchema getSchema() {
      Descriptor desc = Protobufs.getMessageDescriptor(typeRef.getRawClass());
      return new ResourceSchema(new ProtobufToPig().toSchema(desc));
    }

    @Override
    public Tuple toTuple(Object obj) {
      return new ProtobufTuple((Message)obj);
    }
  }
}

@@ -1463,17 +1463,16 @@
 pleImpl%0A
-%0A
     if (
@@ -1542,39 +1542,32 @@
 ProtobufToTuple(
-typeRef
 );%0A    %7D else if
@@ -1643,23 +1643,16 @@
 ToTuple(
-typeRef
 );%0A    %7D
@@ -2510,17 +2510,16 @@
 :%0A   */%0A
-%0A
   privat
@@ -2643,31 +2643,24 @@
 %0A%0A  private 
-static 
 class Thrift
@@ -2695,205 +2695,67 @@
 e %7B%0A
-    private TypeRef%3CTBase%3C?, ?%3E%3E typeRef;%0A    private ThriftToPig%3CTBase%3C?, ?%3E%3E thriftToPig;%0A%0A    @SuppressWarnings(%22unchecked%22)%0A    ThriftToTuple(TypeRef%3C?%3E typeRef) %7B%0A      this.typeRef = (TypeRef
+%0A    @SuppressWarnings(%22unchecked%22)%0A    private ThriftToPig
 %3CTBa
@@ -2767,30 +2767,9 @@
  ?%3E%3E
-)
  
-typeRef;%0A      this.
 thri
@@ -2802,21 +2802,25 @@
 nstance(
-this.
+(TypeRef)
 typeRef)
@@ -2820,22 +2820,16 @@
 ypeRef);
-%0A    %7D
 %0A%0A    @O
@@ -3071,23 +3071,16 @@
 private 
-static 
 class Pr
@@ -3121,192 +3121,53 @@
 e %7B%0A
+%0A
     
-private TypeRef%3CMessage%3E typeRef;%0A%0A    @SuppressWarnings(%22unchecked%22)%0A    ProtobufToTuple(TypeRef%3C?%3E typeRef) %7B%0A      this.typeRef = (TypeRef%3CMessage%3E)typeRef;%0A    %7D%0A%0A    @Override
+@Override @SuppressWarnings(%22unchecked%22)
 %0A   
@@ -3258,16 +3258,32 @@
 criptor(
+(Class%3CMessage%3E)
 typeRef.

cut some more fat