package com.altamiracorp.securegraph.accumulo.mr;

import com.altamiracorp.securegraph.accumulo.AccumuloConstants;
import org.apache.accumulo.core.client.mapreduce.AccumuloRowInputFormat;
import org.apache.accumulo.core.client.security.tokens.PasswordToken;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.security.Authorizations;
import org.apache.accumulo.core.util.CachedConfiguration;
import org.apache.accumulo.core.util.PeekingIterator;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.Mapper;
import org.apache.hadoop.mapreduce.Reducer;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;

public class ItemCountMR extends Configured implements Tool {
    private static final IntWritable INT_WRITABLE_1 = new IntWritable(1);
    private static final Text DATA_COUNT = new Text(AccumuloConstants.DATA_ROW_KEY_PREFIX);
    private static final Text VERTEX_COUNT = new Text(AccumuloConstants.VERTEX_ROW_KEY_PREFIX);
    private static final Text EDGE_COUNT = new Text(AccumuloConstants.EDGE_ROW_KEY_PREFIX);

    public static class CountMapper extends Mapper<Text, PeekingIterator<Map.Entry<Key, Value>>, Text, IntWritable> {
        public void map(Key key, PeekingIterator<Map.Entry<Key, Value>> row, Context context) throws IOException, InterruptedException {
            if (row.hasNext()) {
                Map.Entry<Key, Value> col = row.next();
                String prefix = "" + col.getKey().getRow().toString().charAt(0);
                if (AccumuloConstants.DATA_ROW_KEY_PREFIX.equals(prefix)) {
                    context.write(DATA_COUNT, INT_WRITABLE_1);
                    return;
                }
                if (AccumuloConstants.VERTEX_ROW_KEY_PREFIX.equals(prefix)) {
                    context.write(VERTEX_COUNT, INT_WRITABLE_1);
                    return;
                }
                if (AccumuloConstants.EDGE_ROW_KEY_PREFIX.equals(prefix)) {
                    context.write(EDGE_COUNT, INT_WRITABLE_1);
                    return;
                }
            }
        }
    }

    public static class CountReducer extends Reducer<Text, IntWritable, Text, IntWritable> {
        public void reduce(Text key, Iterator<IntWritable> values, Context context) throws IOException, InterruptedException {
            int sum = 0;
            while (values.hasNext()) {
                sum += values.next().get();
            }
            context.write(key, new IntWritable(sum));
        }
    }


    @Override
    public int run(String[] args) throws Exception {
        if (args.length != 6) {
            throw new IllegalArgumentException("Usage : " + ItemCountMR.class.getSimpleName() + " <instance name> <zookeepers> <user> <password> <table> <auths>");
        }
        String instanceName = args[0];
        String zookeepers = args[1];
        String user = args[2];
        String password = args[3];
        String table = args[4];
        String authorizations = args[5];

        String jobName = this.getClass().getSimpleName() + "_" + System.currentTimeMillis();
        String outputDirectory = "/tmp/" + jobName;

        Job job = Job.getInstance(getConf(), jobName);
        job.setJarByClass(this.getClass());

        job.setInputFormatClass(AccumuloRowInputFormat.class);
        String[] authorizationsArray = authorizations.split(",");
        if (authorizationsArray.length == 1 && authorizationsArray[0].length() == 0) {
            authorizationsArray = new String[0];
        }
        AccumuloRowInputFormat.setConnectorInfo(job, user, new PasswordToken(password.getBytes()));
        AccumuloRowInputFormat.setInputTableName(job, table);
        AccumuloRowInputFormat.setScanAuthorizations(job, new Authorizations(authorizationsArray));
        AccumuloRowInputFormat.setZooKeeperInstance(job, instanceName, zookeepers);

        job.setMapperClass(CountMapper.class);
        job.setMapOutputKeyClass(Text.class);
        job.setMapOutputValueClass(IntWritable.class);

        job.setCombinerClass(CountReducer.class);
        job.setReducerClass(CountReducer.class);

        job.setOutputFormatClass(TextOutputFormat.class);
        TextOutputFormat.setOutputPath(job, new Path(outputDirectory));

        job.waitForCompletion(true);

        return job.isSuccessful() ? 0 : 1;
    }


    public static void main(String[] args) throws Exception {
        int res = ToolRunner.run(CachedConfiguration.getInstance(), new ItemCountMR(), args);
        System.exit(res);
    }
}


@@ -925,16 +925,73 @@
 lRunner;
+%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;
 %0A%0Aimport
@@ -1124,16 +1124,101 @@
  Tool %7B%0A
+    private static final Logger LOGGER = LoggerFactory.getLogger(ItemCountMR.class);%0A
     priv
@@ -1823,101 +1823,8 @@
 n %7B%0A
-            if (row.hasNext()) %7B%0A                Map.Entry%3CKey, Value%3E col = row.next();%0A    
     
@@ -1856,20 +1856,11 @@
 %22 + 
-col.getKey()
+key
 .get
@@ -1887,20 +1887,16 @@
 rAt(0);%0A
-    
         
@@ -1967,36 +1967,32 @@
                 
-    
 context.write(DA
@@ -2026,36 +2026,32 @@
                 
-   
- 
 return;%0A        
@@ -2050,34 +2050,26 @@
             
-    %7D%0A    
+%7D%0A
             
@@ -2122,36 +2122,32 @@
 uals(prefix)) %7B%0A
-    
                 
@@ -2199,36 +2199,32 @@
                 
-    
 return;%0A        
@@ -2223,34 +2223,26 @@
             
-    %7D%0A    
+%7D%0A
             
@@ -2309,36 +2309,32 @@
                 
- 
-   
 context.write(ED
@@ -2368,36 +2368,32 @@
                 
-    
 return;%0A        
@@ -2398,29 +2398,78 @@
     
+  %7D%0A
       
-%7D%0A
       
-      %7D
+LOGGER.warn(%22Unrecognized row key prefix: %22 + prefix);
 %0A   

use key