package com.squareup.cascading2.function;

import cascading.operation.BaseOperation;
import cascading.operation.Function;
import cascading.tuple.Fields;
import com.google.protobuf.Descriptors;
import com.google.protobuf.Message;
import com.squareup.cascading2.util.Util;
import com.squareup.cascading_helpers.operation.KnowsEmittedClasses;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public abstract class AbstractExpandProto<T extends Message> extends BaseOperation implements Function, KnowsEmittedClasses {
  protected final String messageClassName;
  protected final String[] fieldsToExtract;
  protected transient Descriptors.FieldDescriptor[] fieldDescriptorsToExtract;

  protected AbstractExpandProto(Class<T> messageClass, Fields fieldDeclaration, String[] fieldsToExtract) {
    super(1, fieldDeclaration);
    this.messageClassName = messageClass.getName();
    this.fieldsToExtract = fieldsToExtract;
  }

  protected static <T extends Message> String[] getAllFields(Class<T> messageClass) {
    try {
      Method m = messageClass.getMethod("newBuilder");
      Message.Builder builder = (Message.Builder) m.invoke(new Object[]{});

      List<String> fieldNames = new ArrayList<String>();
      for (Descriptors.FieldDescriptor fieldDesc : builder.getDescriptorForType().getFields()) {
        fieldNames.add(fieldDesc.getName());
      }
      return fieldNames.toArray(new String[fieldNames.size()]);
    } catch (IllegalAccessException e) {
      throw new RuntimeException(e);
    } catch (NoSuchMethodException e) {
      throw new RuntimeException(e);
    } catch (InvocationTargetException e) {
      throw new RuntimeException(e);
    }
  }

  protected Descriptors.FieldDescriptor[] getFieldDescriptorsToExtract() {
    if (fieldDescriptorsToExtract == null) {
      Message.Builder builder = Util.builderFromMessageClass(messageClassName);

      List <Descriptors.FieldDescriptor> fieldDescriptors = new ArrayList<Descriptors.FieldDescriptor>();
      for (int i = 0; i < fieldsToExtract.length; i++) {
        fieldDescriptors.add(builder.getDescriptorForType().findFieldByName(fieldsToExtract[i]));
      }

      fieldDescriptorsToExtract = fieldDescriptors.toArray(new Descriptors.FieldDescriptor[fieldDescriptors.size()]);
    }

    return fieldDescriptorsToExtract;
  }

  @Override
  public final Set<Class> getEmittedClasses() {
    Set<Class> results = new HashSet<Class>();
    for (Descriptors.FieldDescriptor fieldDesc : getFieldDescriptorsToExtract()) {
      Descriptors.FieldDescriptor.Type type = fieldDesc.getType();
      if (type.getJavaType() == Descriptors.FieldDescriptor.JavaType.MESSAGE) {
        results.add(Util.messageClassFromFieldDesc(messageClassName, fieldDesc));
      }
    }

    return results;
  }
}

@@ -1147,75 +1147,8 @@
     
-try %7B%0A      Method m = messageClass.getMethod(%22newBuilder%22);%0A      
 Mess
@@ -1163,65 +1163,73 @@
 der 
+builder = Util.
 builder
- = (
+From
 Message
-.Builder) m.invoke(new Object%5B%5D%7B%7D
+Class(messageClass.getName()
 );%0A%0A
-  
     
@@ -1275,26 +1275,24 @@
 %3CString%3E();%0A
-  
     for (Des
@@ -1370,26 +1370,24 @@
 Fields()) %7B%0A
-  
       fieldN
@@ -1421,22 +1421,18 @@
 ));%0A    
-  %7D%0A  
+%7D%0A
     retu
@@ -1489,250 +1489,8 @@
 %5D);%0A
-    %7D catch (IllegalAccessException e) %7B%0A      throw new RuntimeException(e);%0A    %7D catch (NoSuchMethodException e) %7B%0A      throw new RuntimeException(e);%0A    %7D catch (InvocationTargetException e) %7B%0A      throw new RuntimeException(e);%0A    %7D%0A
   %7D%0A
@@ -1812,22 +1812,26 @@
 or (
-int i = 0; i %3C
+String fieldName :
  fie
@@ -1846,20 +1846,8 @@
 ract
-.length; i++
 ) %7B%0A
@@ -1931,21 +1931,12 @@
 ield
-sToExtract%5Bi%5D
+Name
 ));%0A

address a few minor nits