/*
 * 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.pig.piggybank.storage;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.io.compress.BZip2Codec;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskID;
import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.pig.StoreFunc;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MRConfiguration;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.util.StorageUtil;
import org.apache.xml.utils.StringBufferPool;

import com.google.common.base.Strings;

/**
 * The UDF is useful for splitting the output data into a bunch of directories
 * and files dynamically based on user specified key field in the output tuple.
 * 
 * Sample usage: <code>
 * A = LOAD 'mydata' USING PigStorage() as (a, b, c);
 * STORE A INTO '/my/home/output' USING MultiStorage('/my/home/output','0', 'bz2', '\\t');
 * </code> Parameter details:- ========== <b>/my/home/output </b>(Required) :
 * The DFS path where output directories and files will be created. <b> 0
 * </b>(Required) : Index of field whose values should be used to create
 * directories and files( field 'a' in this case). <b>'bz2' </b>(Optional) : The
 * compression type. Default is 'none'. Supported types are:- 'none', 'gz' and
 * 'bz2' <b> '\\t' </b>(Optional) : Output field separator.
 * 
 * Let 'a1', 'a2' be the unique values of field 'a'. Then output may look like
 * this
 * 
 * /my/home/output/a1/a1-0000 /my/home/output/a1/a1-0001
 * /my/home/output/a1/a1-0002 ... /my/home/output/a2/a2-0000
 * /my/home/output/a2/a2-0001 /my/home/output/a2/a2-0002
 * 
 * The prefix '0000*' is the task-id of the mapper/reducer task executing this
 * store. In case user does a GROUP BY on the field followed by MultiStorage(),
 * then its imperative that all tuples for a particular group will go exactly to
 * 1 reducer. So in the above case for e.g. there will be only 1 file each under
 * 'a1' and 'a2' directories.
 * 
 * If the output is compressed,then the sub directories and the output files will
 * be having the extension. Say for example in the above case if bz2 is used one file 
 * will look like ;/my/home/output.bz2/a1.bz2/a1-0000.bz2
 *
 * Key field can also be a comma separated list of indices e.g. '0,1' - in this case
 * storage will be multi-level:
 * /my/home/output/a1/b1/a1-b1-0000
 * /my/home/output/a1/b2/a1-b2-0000
 * There is also an option to leave key values out of storage, see isRemoveKeys.
 */
public class MultiStorage extends StoreFunc {

  private static final String KEYFIELD_DELIMETER = ",";
  private Path outputPath; // User specified output Path
  private final List<Integer> splitFieldIndices= new ArrayList<Integer>(); // Indices of the key fields
  private String fieldDel; // delimiter of the output record.
  private Compression comp; // Compression type of output data.
  private boolean isRemoveKeys = false;
  
  // Compression types supported by this store
  enum Compression {
    none, bz2, bz, gz;
  };

  public MultiStorage(String parentPathStr, String splitFieldIndex) {
    this(parentPathStr, splitFieldIndex, "none");
  }

  public MultiStorage(String parentPathStr, String splitFieldIndex,
      String compression) {
    this(parentPathStr, splitFieldIndex, compression, "\\t");
  }

  public MultiStorage(String parentPathStr, String splitFieldIndex,
      String compression, String fieldDel) {
    this(parentPathStr, splitFieldIndex, compression, fieldDel, "false");
  }

  /**
   * Constructor
   *
   * @param parentPathStr
   *          Parent output dir path (this will be specified in store statement,
   *            so MultiStorage don't use this parameter in reality. However, we don't
   *            want to change the construct to break backward compatibility)
   * @param splitFieldIndex
   *          key field index
   * @param compression
   *          'bz2', 'bz', 'gz' or 'none'
   * @param fieldDel
   *          Output record field delimiter.
   * @param isRemoveKeys
   *          Removes key columns from result during write.
   */
  public MultiStorage(String parentPathStr, String splitFieldIndex,
                      String compression, String fieldDel, String isRemoveKeys) {
    this.isRemoveKeys = Boolean.parseBoolean(isRemoveKeys);
    this.outputPath = new Path(parentPathStr);

    String[] splitFieldIndices = splitFieldIndex.split(KEYFIELD_DELIMETER);
    for (String splitFieldIndexString : splitFieldIndices){
      this.splitFieldIndices.add(Integer.parseInt(splitFieldIndexString));
    }

    this.fieldDel = fieldDel;
    try {
      this.comp = (compression == null) ? Compression.none : Compression
              .valueOf(compression.toLowerCase());
    } catch (IllegalArgumentException e) {
      System.err.println("Exception when converting compression string: "
              + compression + " to enum. No compression will be used");
      this.comp = Compression.none;
    }
  }

  //--------------------------------------------------------------------------
  // Implementation of StoreFunc

  private RecordWriter<List<String>, Tuple> writer;
  
  @Override
  public void putNext(Tuple tuple) throws IOException {
    for (int splitFieldIndex : this.splitFieldIndices) {
      if (tuple.size() <= splitFieldIndex) {
        throw new IOException("split field index:" + splitFieldIndex
                + " >= tuple size:" + tuple.size());
      }
    }
    List<String> fields = new ArrayList<String>();
    for (int splitFieldIndex : this.splitFieldIndices){
      try {
        fields.add(String.valueOf(tuple.get(splitFieldIndex)));
      } catch (ExecException exec) {
        throw new IOException(exec);
      }
    }
    try {
      writer.write(fields, tuple);
    } catch (InterruptedException e) {
      throw new IOException(e);
    }
  }
  
  @SuppressWarnings("unchecked")
  @Override
  public OutputFormat getOutputFormat() throws IOException {
      MultiStorageOutputFormat format = new MultiStorageOutputFormat();
      format.setKeyValueSeparator(fieldDel);
      if (this.isRemoveKeys){
        format.setSkipIndices(this.splitFieldIndices);
      }
      return format;
  }
    
  @SuppressWarnings("unchecked")
  @Override
  public void prepareToWrite(RecordWriter writer) throws IOException {
      this.writer = writer;
  }
    
  @Override
  public void setStoreLocation(String location, Job job) throws IOException {
    job.getConfiguration().set(MRConfiguration.TEXTOUTPUTFORMAT_SEPARATOR, "");
    FileOutputFormat.setOutputPath(job, new Path(location));
    if (comp == Compression.bz2 || comp == Compression.bz) {
      FileOutputFormat.setCompressOutput(job, true);
      FileOutputFormat.setOutputCompressorClass(job,  BZip2Codec.class);
    } else if (comp == Compression.gz) {
      FileOutputFormat.setCompressOutput(job, true);
      FileOutputFormat.setOutputCompressorClass(job, GzipCodec.class);
    }
  }

  @Override
  public Boolean supportsParallelWriteToStoreLocation() {
    return false;
  }

  //--------------------------------------------------------------------------
  // Implementation of OutputFormat
  
  public static class MultiStorageOutputFormat extends
  TextOutputFormat<List<String>, Tuple> {

    private String keyValueSeparator = "\\t";
    private byte fieldDel = '\t';
    private List<Integer> skipIndices = null;

    @Override
    public RecordWriter<List<String>, Tuple>
    getRecordWriter(TaskAttemptContext context
                ) throws IOException, InterruptedException {
    
      final TaskAttemptContext ctx = context;
        
      return new RecordWriter<List<String>, Tuple>() {

        private Map<List<String>, MyLineRecordWriter> storeMap =
              new HashMap<List<String>, MyLineRecordWriter>();
          
        private static final int BUFFER_SIZE = 1024;
          
        private ByteArrayOutputStream mOut = 
              new ByteArrayOutputStream(BUFFER_SIZE);
                           
        @Override
        public void write(List<String> key, Tuple val) throws IOException {
          int sz = val.size();
          for (int i = 0; i < sz; i++) {
            Object field;
            try {
              field = val.get(i);
            } catch (ExecException ee) {
              throw ee;
            }

            boolean skipCurrentField = skipIndices != null && skipIndices.contains(i);

            if (!skipCurrentField) {
              StorageUtil.putField(mOut, field);
            }

            if (i != sz - 1 && !skipCurrentField) {
              mOut.write(fieldDel);
            }
          }
              
          getStore(key).write(null, new Text(mOut.toByteArray()));

          mOut.reset();
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException { 
          for (MyLineRecordWriter out : storeMap.values()) {
            out.close(context);
          }
        }
      
        private MyLineRecordWriter getStore(List<String> fieldValues) throws IOException {
          MyLineRecordWriter store = storeMap.get(fieldValues);
          if (store == null) {                  
            DataOutputStream os = createOutputStream(fieldValues);
            store = new MyLineRecordWriter(os, keyValueSeparator);
            storeMap.put(fieldValues, store);
          }
          return store;
        }
          
        private DataOutputStream createOutputStream(List<String> fieldValues) throws IOException {
          Configuration conf = ctx.getConfiguration();
          TaskID taskId = ctx.getTaskAttemptID().getTaskID();
          
          // Check whether compression is enabled, if so get the extension and add them to the path
          boolean isCompressed = getCompressOutput(ctx);
          CompressionCodec codec = null;
          String extension = "";
          if (isCompressed) {
             Class<? extends CompressionCodec> codecClass = 
                getOutputCompressorClass(ctx, GzipCodec.class);
             codec = (CompressionCodec) ReflectionUtils.newInstance(codecClass, ctx.getConfiguration());
             extension = codec.getDefaultExtension();
          }

          NumberFormat nf = NumberFormat.getInstance();
          nf.setMinimumIntegerDigits(4);

          StringBuffer pathStringBuffer = new StringBuffer();
          for (String fieldValue : fieldValues){
            String safeFieldValue = fieldValue.replaceAll("\\/","-");
            pathStringBuffer.append(safeFieldValue);
            pathStringBuffer.append("/");
          }
          pathStringBuffer.deleteCharAt(pathStringBuffer.length()-1);
          String pathString = pathStringBuffer.toString();
          String idString = pathString.replaceAll("\\/","-");

          if (!Strings.isNullOrEmpty(extension)){
            pathString = pathString.replaceAll("\\/",extension+"\\/");
          }

          Path path = new Path(pathString+extension, idString + '-'
                + nf.format(taskId.getId())+extension);
          Path workOutputPath = ((FileOutputCommitter)getOutputCommitter(ctx)).getWorkPath();
          Path file = new Path(workOutputPath, path);
          FileSystem fs = file.getFileSystem(conf);                
          FSDataOutputStream fileOut = fs.create(file, false);
          
          if (isCompressed)
             return new DataOutputStream(codec.createOutputStream(fileOut));
          else
             return fileOut;
        }
          
      };
    }
  
    public void setKeyValueSeparator(String sep) {
      keyValueSeparator = sep;
      fieldDel = StorageUtil.parseFieldDel(keyValueSeparator);  
    }

    public void setSkipIndices(List<Integer> skipIndices) {
      this.skipIndices = skipIndices;
    }

    //------------------------------------------------------------------------
  //
  
    protected static class MyLineRecordWriter
    extends TextOutputFormat.LineRecordWriter<WritableComparable, Text> {

      public MyLineRecordWriter(DataOutputStream out, String keyValueSeparator) {
        super(out, keyValueSeparator);
      }
    }
  }

}
