/*
    * Licensed to the Apache Software Foundation (ASF) under one
    * or more contributor license agreements.  See the NOTICE file
    * distributed with this work for additional information
    * regarding copyright ownership.  The ASF licenses this file
    * to you under the Apache License, Version 2.0 (the
    * "License"); you may not use this file except in compliance
    * with the License.  You may obtain a copy of the License at
    *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.giraph.io.hcatalog;
  
  import org.apache.commons.cli.CommandLine;
  import org.apache.commons.cli.CommandLineParser;
  import org.apache.commons.cli.GnuParser;
  import org.apache.commons.cli.HelpFormatter;
  import org.apache.commons.cli.Options;
  import org.apache.commons.cli.ParseException;
  import org.apache.giraph.graph.Computation;
  import org.apache.giraph.io.EdgeInputFormat;
  import org.apache.giraph.io.VertexInputFormat;
  import org.apache.giraph.io.VertexOutputFormat;
  import org.apache.giraph.job.GiraphJob;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.hive.conf.HiveConf;
  import org.apache.hadoop.util.Tool;
  import org.apache.hadoop.util.ToolRunner;
  import org.apache.hcatalog.mapreduce.HCatOutputFormat;
  import org.apache.hcatalog.mapreduce.InputJobInfo;
  import org.apache.hcatalog.mapreduce.OutputJobInfo;
  import org.apache.log4j.Logger;
  
  import com.google.common.collect.Lists;
  
  import java.io.File;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  
  /**
   * Hive Giraph Runner
   */
  public class HCatGiraphRunner implements Tool {
    /**
     * logger
     */
    private static final Logger LOG = Logger.getLogger(HCatGiraphRunner.class);
    /**
     * workers
     */
    protected int workers;
    /**
     * is verbose
     */
    protected boolean isVerbose;
    /**
     * output table partitions
     */
    protected Map<String, String> outputTablePartitionValues;
    /**
     * dbName
     */
    protected String dbName;
    /**
     * vertex input table name
     */
    protected String vertexInputTableName;
    /**
     * vertex input table filter
     */
    protected String vertexInputTableFilterExpr;
    /**
     * edge input table name
     */
    protected String edgeInputTableName;
    /**
     * edge input table filter
     */
    protected String edgeInputTableFilterExpr;
    /**
     * output table name
     */
    protected String outputTableName;
    /** Configuration */
    private Configuration conf;
    /** Skip output? (Useful for testing without writing) */
    private boolean skipOutput = false;
  
    /**
    * computation class.
   */
   private Class<? extends Computation> computationClass;
   /**
    * vertex input format internal.
    */
   private Class<? extends VertexInputFormat> vertexInputFormatClass;
   /**
    * edge input format internal.
    */
   private Class<? extends EdgeInputFormat> edgeInputFormatClass;
   /**
   * vertex output format internal.
   */
   private Class<? extends VertexOutputFormat> vertexOutputFormatClass;
 
   /**
   * Giraph runner class.
    *
   * @param computationClass Computation class
   * @param vertexInputFormatClass Vertex input format
   * @param edgeInputFormatClass Edge input format
   * @param vertexOutputFormatClass Output format
   */
   protected HCatGiraphRunner(
       Class<? extends Computation> computationClass,
       Class<? extends VertexInputFormat> vertexInputFormatClass,
       Class<? extends EdgeInputFormat> edgeInputFormatClass,
       Class<? extends VertexOutputFormat> vertexOutputFormatClass) {
     this.computationClass = computationClass;
     this.vertexInputFormatClass = vertexInputFormatClass;
     this.edgeInputFormatClass = edgeInputFormatClass;
     this.vertexOutputFormatClass = vertexOutputFormatClass;
     this.conf = new HiveConf(getClass());
   }
 
   /**
   * main method
   * @param args system arguments
   * @throws Exception any errors from Hive Giraph Runner
   */
   public static void main(String[] args) throws Exception {
     System.exit(ToolRunner.run(
         new HCatGiraphRunner(null, null, null, null), args));
   }
 
   @Override
   public final int run(String[] args) throws Exception {
     // process args
     try {
       processArguments(args);
     } catch (InterruptedException e) {
       return 0;
     } catch (IllegalArgumentException e) {
       System.err.println(e.getMessage());
       return -1;
     }
 
     // additional configuration for Hive
     adjustConfigurationForHive(getConf());
 
     // setup GiraphJob
     GiraphJob job = new GiraphJob(getConf(), getClass().getName());
     job.getConfiguration().setComputationClass(computationClass);
 
     // setup input from Hive
     if (vertexInputFormatClass != null) {
       InputJobInfo vertexInputJobInfo = InputJobInfo.create(dbName,
           vertexInputTableName, vertexInputTableFilterExpr);
       GiraphHCatInputFormat.setVertexInput(job.getInternalJob(),
           vertexInputJobInfo);
       job.getConfiguration().setVertexInputFormatClass(vertexInputFormatClass);
     }
     if (edgeInputFormatClass != null) {
       InputJobInfo edgeInputJobInfo = InputJobInfo.create(dbName,
           edgeInputTableName, edgeInputTableFilterExpr);
       GiraphHCatInputFormat.setEdgeInput(job.getInternalJob(),
           edgeInputJobInfo);
       job.getConfiguration().setEdgeInputFormatClass(edgeInputFormatClass);
     }
 
     // setup output to Hive
     HCatOutputFormat.setOutput(job.getInternalJob(), OutputJobInfo.create(
         dbName, outputTableName, outputTablePartitionValues));
     HCatOutputFormat.setSchema(job.getInternalJob(),
         HCatOutputFormat.getTableSchema(job.getInternalJob()));
     if (skipOutput) {
       LOG.warn("run: Warning - Output will be skipped!");
     } else {
       job.getConfiguration().setVertexOutputFormatClass(
           vertexOutputFormatClass);
     }
 
     job.getConfiguration().setWorkerConfiguration(workers, workers, 100.0f);
     initGiraphJob(job);
 
     return job.run(isVerbose) ? 0 : -1;
   }
 
   /**
   * set hive configuration
   * @param conf Configuration argument
   */
   private static void adjustConfigurationForHive(Configuration conf) {
     // when output partitions are used, workers register them to the
     // metastore at cleanup stage, and on HiveConf's initialization, it
     // looks for hive-site.xml from.
     addToStringCollection(conf, "tmpfiles", conf.getClassLoader()
         .getResource("hive-site.xml").toString());
 
     // Also, you need hive.aux.jars as well
     // addToStringCollection(conf, "tmpjars",
     // conf.getStringCollection("hive.aux.jars.path"));
 
     // Or, more effectively, we can provide all the jars client needed to
     // the workers as well
     String[] hadoopJars = System.getenv("HADOOP_CLASSPATH").split(
         File.pathSeparator);
     List<String> hadoopJarURLs = Lists.newArrayList();
     for (String jarPath : hadoopJars) {
       File file = new File(jarPath);
       if (file.exists() && file.isFile()) {
         String jarURL = file.toURI().toString();
         hadoopJarURLs.add(jarURL);
       }
     }
     addToStringCollection(conf, "tmpjars", hadoopJarURLs);
   }
 
   /**
   * process arguments
   * @param args to process
   * @return CommandLine instance
   * @throws ParseException error parsing arguments
   * @throws InterruptedException interrupted
   */
   private CommandLine processArguments(String[] args) throws ParseException,
             InterruptedException {
     Options options = new Options();
     options.addOption("h", "help", false, "Help");
     options.addOption("v", "verbose", false, "Verbose");
     options.addOption("D", "hiveconf", true,
                 "property=value for Hive/Hadoop configuration");
     options.addOption("w", "workers", true, "Number of workers");
     if (computationClass == null) {
       options.addOption(null, "computationClass", true,
           "Giraph Computation class to use");
     }
     if (vertexInputFormatClass == null) {
       options.addOption(null, "vertexInputFormatClass", true,
           "Giraph HCatalogVertexInputFormat class to use");
     }
     if (edgeInputFormatClass == null) {
       options.addOption(null, "edgeInputFormatClass", true,
           "Giraph HCatalogEdgeInputFormat class to use");
     }
 
     if (vertexOutputFormatClass == null) {
       options.addOption(null, "vertexOutputFormatClass", true,
           "Giraph HCatalogVertexOutputFormat class to use");
     }
 
     options.addOption("db", "dbName", true, "Hive database name");
     options.addOption("vi", "vertexInputTable", true,
         "Vertex input table name");
     options.addOption("VI", "vertexInputFilter", true,
         "Vertex input table filter expression (e.g., \"a<2 AND b='two'\"");
     options.addOption("ei", "edgeInputTable", true,
         "Edge input table name");
     options.addOption("EI", "edgeInputFilter", true,
         "Edge input table filter expression (e.g., \"a<2 AND b='two'\"");
     options.addOption("o", "outputTable", true, "Output table name");
     options.addOption("O", "outputPartition", true,
         "Output table partition values (e.g., \"a=1,b=two\")");
     options.addOption("s", "skipOutput", false, "Skip output?");
 
     addMoreOptions(options);
 
     CommandLineParser parser = new GnuParser();
     final CommandLine cmdln = parser.parse(options, args);
     if (args.length == 0 || cmdln.hasOption("help")) {
       new HelpFormatter().printHelp(getClass().getName(), options, true);
       throw new InterruptedException();
     }
 
     // Giraph classes
     if (cmdln.hasOption("computationClass")) {
       computationClass = findClass(cmdln.getOptionValue("computationClass"),
           Computation.class);
     }
     if (cmdln.hasOption("vertexInputFormatClass")) {
       vertexInputFormatClass = findClass(
           cmdln.getOptionValue("vertexInputFormatClass"),
           HCatalogVertexInputFormat.class);
     }
     if (cmdln.hasOption("edgeInputFormatClass")) {
       edgeInputFormatClass = findClass(
           cmdln.getOptionValue("edgeInputFormatClass"),
           HCatalogEdgeInputFormat.class);
     }
 
     if (cmdln.hasOption("vertexOutputFormatClass")) {
       vertexOutputFormatClass = findClass(
           cmdln.getOptionValue("vertexOutputFormatClass"),
           HCatalogVertexOutputFormat.class);
     }
 
     if (cmdln.hasOption("skipOutput")) {
       skipOutput = true;
     }
 
     if (computationClass == null) {
       throw new IllegalArgumentException(
           "Need the Giraph Computation class name (-computationClass) to use");
     }
     if (vertexInputFormatClass == null && edgeInputFormatClass == null) {
       throw new IllegalArgumentException(
           "Need at least one of Giraph VertexInputFormat " +
               "class name (-vertexInputFormatClass) and " +
               "EdgeInputFormat class name (-edgeInputFormatClass)");
     }
     if (vertexOutputFormatClass == null) {
       throw new IllegalArgumentException(
           "Need the Giraph VertexOutputFormat " +
               "class name (-vertexOutputFormatClass) to use");
     }
     if (!cmdln.hasOption("workers")) {
       throw new IllegalArgumentException(
           "Need to choose the number of workers (-w)");
     }
     if (!cmdln.hasOption("vertexInputTable") &&
         vertexInputFormatClass != null) {
       throw new IllegalArgumentException(
           "Need to set the vertex input table name (-vi)");
     }
     if (!cmdln.hasOption("edgeInputTable") &&
         edgeInputFormatClass != null) {
       throw new IllegalArgumentException(
           "Need to set the edge input table name (-ei)");
     }
     if (!cmdln.hasOption("outputTable")) {
       throw new IllegalArgumentException(
           "Need to set the output table name (-o)");
     }
     dbName = cmdln.getOptionValue("dbName", "default");
     vertexInputTableName = cmdln.getOptionValue("vertexInputTable");
     vertexInputTableFilterExpr = cmdln.getOptionValue("vertexInputFilter");
     edgeInputTableName = cmdln.getOptionValue("edgeInputTable");
     edgeInputTableFilterExpr = cmdln.getOptionValue("edgeInputFilter");
     outputTableName = cmdln.getOptionValue("outputTable");
     outputTablePartitionValues = HiveUtils.parsePartitionValues(cmdln
                 .getOptionValue("outputPartition"));
     workers = Integer.parseInt(cmdln.getOptionValue("workers"));
     isVerbose = cmdln.hasOption("verbose");
 
     // pick up -hiveconf arguments
     for (String hiveconf : cmdln.getOptionValues("hiveconf")) {
       String[] keyval = hiveconf.split("=", 2);
       if (keyval.length == 2) {
         String name = keyval[0];
         String value = keyval[1];
         if (name.equals("tmpjars") || name.equals("tmpfiles")) {
           addToStringCollection(
                   conf, name, value);
         } else {
           conf.set(name, value);
         }
       }
     }
 
     processMoreArguments(cmdln);
 
     return cmdln;
   }
 
   /**
   * add string to collection
   * @param conf Configuration
   * @param name name to add
   * @param values values for collection
   */
   private static void addToStringCollection(Configuration conf, String name,
                                               String... values) {
     addToStringCollection(conf, name, Arrays.asList(values));
   }
 
   /**
   * add string to collection
   * @param conf Configuration
   * @param name to add
   * @param values values for collection
   */
   private static void addToStringCollection(
           Configuration conf, String name, Collection
           <? extends String> values) {
     Collection<String> tmpfiles = conf.getStringCollection(name);
     tmpfiles.addAll(values);
     conf.setStrings(name, tmpfiles.toArray(new String[tmpfiles.size()]));
   }
 
   /**
   *
   * @param className to find
   * @param base  base class
   * @param <T> class type found
   * @return type found
   */
   private <T> Class<? extends T> findClass(String className, Class<T> base) {
     try {
       Class<?> cls = Class.forName(className);
       if (base.isAssignableFrom(cls)) {
         return cls.asSubclass(base);
       }
       return null;
     } catch (ClassNotFoundException e) {
       throw new IllegalArgumentException(className + ": Invalid class name");
     }
   }
 
   @Override
   public final Configuration getConf() {
     return conf;
   }
 
   @Override
   public final void setConf(Configuration conf) {
     this.conf = conf;
   }
 
   /**
   * Override this method to add more command-line options. You can process
   * them by also overriding {@link #processMoreArguments(CommandLine)}.
   *
   * @param options Options
   */
   protected void addMoreOptions(Options options) {
   }
 
   /**
   * Override this method to process additional command-line arguments. You
   * may want to declare additional options by also overriding
   * {@link #addMoreOptions(Options)}.
   *
   * @param cmd Command
   */
   protected void processMoreArguments(CommandLine cmd) {
   }
 
   /**
   * Override this method to do additional setup with the GiraphJob that will
   * run.
   *
   * @param job
   *            GiraphJob that is going to run
   */
   protected void initGiraphJob(GiraphJob job) {
     LOG.info(getClass().getSimpleName() + " with");
     String prefix = "\t";
     LOG.info(prefix + "-computationClass=" +
          computationClass.getCanonicalName());
     if (vertexInputFormatClass != null) {
       LOG.info(prefix + "-vertexInputFormatClass=" +
           vertexInputFormatClass.getCanonicalName());
     }
     if (edgeInputFormatClass != null) {
       LOG.info(prefix + "-edgeInputFormatClass=" +
           edgeInputFormatClass.getCanonicalName());
     }
     LOG.info(prefix + "-vertexOutputFormatClass=" +
         vertexOutputFormatClass.getCanonicalName());
     if (vertexInputTableName != null) {
       LOG.info(prefix + "-vertexInputTable=" + vertexInputTableName);
     }
     if (vertexInputTableFilterExpr != null) {
       LOG.info(prefix + "-vertexInputFilter=\"" +
           vertexInputTableFilterExpr + "\"");
     }
     if (edgeInputTableName != null) {
       LOG.info(prefix + "-edgeInputTable=" + edgeInputTableName);
     }
     if (edgeInputTableFilterExpr != null) {
       LOG.info(prefix + "-edgeInputFilter=\"" +
           edgeInputTableFilterExpr + "\"");
     }
     LOG.info(prefix + "-outputTable=" + outputTableName);
     if (outputTablePartitionValues != null) {
       LOG.info(prefix + "-outputPartition=\"" +
           outputTablePartitionValues + "\"");
     }
     LOG.info(prefix + "-workers=" + workers);
   }
 }