package org.calrissian.accumulorecipes.eventstore.pig;

import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import groovy.lang.Binding;
import groovy.lang.GroovyShell;
import org.apache.accumulo.core.client.AccumuloException;
import org.apache.accumulo.core.client.AccumuloSecurityException;
import org.apache.accumulo.core.client.TableNotFoundException;
import org.apache.accumulo.core.security.Authorizations;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.pig.LoadFunc;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigSplit;
import org.apache.pig.data.Tuple;
import org.calrissian.accumulorecipes.eventstore.hadoop.EventInputFormat;
import org.calrissian.mango.criteria.builder.QueryBuilder;
import org.calrissian.mango.criteria.domain.Node;
import org.calrissian.mango.uri.support.UriUtils;
import org.joda.time.DateTime;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Set;

import static com.google.common.collect.Sets.newHashSet;
import static java.util.Arrays.asList;
import static org.apache.commons.lang.StringUtils.splitPreserveAllTokens;
import static org.calrissian.accumulorecipes.eventstore.impl.AccumuloEventStore.DEFAULT_IDX_TABLE_NAME;
import static org.calrissian.accumulorecipes.eventstore.impl.AccumuloEventStore.DEFAULT_SHARD_TABLE_NAME;

public class EventLoader extends LoadFunc {

  private String accumuloUser;
  private String accumuloPass;
  private String accumuloInst;
  private String zookeepers;

  private String shardTable = DEFAULT_SHARD_TABLE_NAME;
  private String indexTable = DEFAULT_IDX_TABLE_NAME;

  private Node query;

  private DateTime startTime;
  private DateTime endTime;

  @Override
  public void setLocation(String uri, Job job) throws IOException {

    Configuration conf = job.getConfiguration();

    String queryPortion = uri.substring(uri.indexOf("?"), uri.length());
    Multimap<String,String> queryParams = UriUtils.splitQuery(queryPortion);

    String accumuloUser = getProp(queryParams, "accumuloUser");
    String accumuloPass = getProp(queryParams, "accumuloPass");
    String accumuloInst = getProp(queryParams, "accumuloInst");
    String zookeepers = getProp(queryParams, "zookeepers");
    String query = getProp(queryParams, "query");
    String shardTable = getProp(queryParams, "shardTable");
    String indexTable = getProp(queryParams, "indexTable");
    String startTime = getProp(queryParams, "startTime");
    String endTime = getProp(queryParams, "endTime");
    String auths = getProp(queryParams, "auths");
    String selectFields = getProp(queryParams, "selectFields");

    Set<String> fields = selectFields != null ? newHashSet(asList(splitPreserveAllTokens(selectFields, ","))) : null;

    DateTime startDT = DateTime.parse(startTime);
    DateTime endDT = DateTime.parse(endTime);


// call groovy expressions from Java code
    Binding binding = new Binding();
    binding.setVariable("q", new QueryBuilder());
    GroovyShell shell = new GroovyShell(binding);
    QueryBuilder qb = (QueryBuilder) shell.evaluate(query);

    EventInputFormat.setZooKeeperInstance(conf, accumuloInst, zookeepers);
    EventInputFormat.setInputInfo(conf, accumuloUser, accumuloPass.getBytes(), new Authorizations(auths.getBytes()));
    try {
      EventInputFormat.setQueryInfo(conf, startDT.toDate(), endDT.toDate(), qb.build(), fields);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  private String getProp(Multimap<String,String> queryParams, String propKey) {
    Collection<String> props = queryParams.get(propKey);
    if(props.size() > 0)
      return props.iterator().next();
    return null;
  }

  @Override
  public InputFormat getInputFormat() throws IOException {
    return new EventInputFormat();
  }

  @Override
  public void prepareToRead(RecordReader recordReader, PigSplit pigSplit) throws IOException {
  }

  @Override
  public Tuple getNext() throws IOException {
    return null;
  }
}

@@ -96,47 +96,8 @@
 ap;%0A
-import com.google.common.collect.Sets;%0A
 impo
@@ -188,161 +188,16 @@
 ore.
-client.AccumuloException;%0Aimport org.apache.accumulo.core.client.AccumuloSecurityException;%0Aimport org.apache.accumulo.core.client.TableNotFoundException
+data.Key
 ;%0Aim
@@ -255,52 +255,8 @@
 ns;%0A
-import org.apache.commons.lang.StringUtils;%0A
 impo
@@ -577,16 +577,89 @@
 .Tuple;%0A
+import org.calrissian.accumulorecipes.commons.hadoop.StoreEntryWritable;%0A
 import o
@@ -787,58 +787,8 @@
 er;%0A
-import org.calrissian.mango.criteria.domain.Node;%0A
 impo
@@ -897,33 +897,8 @@
 on;%0A
-import java.util.Arrays;%0A
 impo
@@ -1119,577 +1119,108 @@
 ns;%0A
-import static org.calrissian.accumulorecipes.eventstore.impl.AccumuloEventStore.DEFAULT_IDX_TABLE_NAME;%0Aimport static org.calrissian.accumulorecipes.eventstore.impl.AccumuloEventStore.DEFAULT_SHARD_TABLE_NAME;%0A%0Apublic class EventLoader extends LoadFunc %7B%0A%0A  private String accumuloUser;%0A  private String accumuloPass;%0A  private String accumuloInst;%0A  private String zookeepers;%0A%0A  private String shardTable = DEFAULT_SHARD_TABLE_NAME;%0A  private String indexTable = DEFAULT_IDX_TABLE_NAME;%0A%0A  private Node query;%0A%0A  private DateTime startTime;%0A  private DateTime endTime
+%0Apublic class EventLoader extends LoadFunc %7B%0A%0A  private RecordReader%3CKey, StoreEntryWritable%3E reader
 ;%0A%0A 
@@ -3422,16 +3422,48 @@
 ption %7B%0A
+    this.reader = recordReader;%0A
   %7D%0A%0A  @
@@ -3517,31 +3517,185 @@
 ption %7B%0A
-    return null;
+%0A    try %7B%0A      if(!reader.nextKeyValue())%0A        return null;%0A      else %7B%0A%0A      %7D%0A%0A%0A    %7D catch (InterruptedException e) %7B%0A      throw new RuntimeException(e);%0A    %7D
 %0A  %7D%0A%7D%0A

Work on event loader pig load func