// BSD License (http://www.galagosearch.org/license)
   
   package org.galagosearch.core.parse;
   
   import java.io.File;
   import java.io.IOException;
   import java.util.ArrayList;
   import java.util.List;
   import org.galagosearch.core.index.VocabularyReader;
  import org.galagosearch.core.index.VocabularyReader.TermSlot;
  import org.galagosearch.core.index.IndexReader;
  import org.galagosearch.tupleflow.ExNihiloSource;
  import org.galagosearch.tupleflow.FileSource;
  import org.galagosearch.tupleflow.IncompatibleProcessorException;
  import org.galagosearch.tupleflow.Linkage;
  import org.galagosearch.tupleflow.OutputClass;
  import org.galagosearch.tupleflow.Parameters.Value;
  import org.galagosearch.tupleflow.Processor;
  import org.galagosearch.tupleflow.Step;
  import org.galagosearch.tupleflow.TupleFlowParameters;
  import org.galagosearch.tupleflow.execution.ErrorHandler;
  import org.galagosearch.tupleflow.execution.Verified;
  import org.galagosearch.core.types.DocumentSplit;
  
  /***
   * From a set of inputs, splits the input into many DocumentSplit records.
   * This will usually be in a stage by itself at the beginning of a Galago pipeline.
   * This is somewhat similar to FileSource, except that it can autodetect file formats.
   * This splitter can detect ARC, TREC, TRECWEB and corpus files.
   * 
   * @author trevor
   */
  
  @Verified
  @OutputClass(className = "org.galagosearch.core.types.DocumentSplit")
  public class DocumentSource implements ExNihiloSource<DocumentSplit> {
      public Processor processor;
      TupleFlowParameters parameters;
      
      public DocumentSource(TupleFlowParameters parameters) {
          this.parameters = parameters;
      }
  
      private String getExtension(String fileName) {
          String[] fields = fileName.split("//.");
          
          // A filename needs to have a period to have an extension.
          if (fields.length <= 1) {
              return "";
          }
          
          // If the last chunk of the filename is gz, we'll ignore it.
          // The second-to-last bit is the type extension (but only if
          // there are at least three parts to the name).
          if (fields[fields.length-1].equals("gz")) {
              if (fields.length > 2) {
                  return fields[fields.length-2];
              } else {
                  return "";
              }
          }
          
          // No 'gz' extension, so just return the last part.
          return fields[fields.length-1];
      }
  
      private void processCorpusFile(String fileName, String fileType) throws IOException {
          // If this is a big file, we'll split it into roughly 100MB pieces.
          long fileLength = new File(fileName).length();
          long chunkSize = 100 * 1024 * 1024;
          
          IndexReader reader = new IndexReader(fileName);
          VocabularyReader vocabulary = reader.getVocabulary();
          List<TermSlot> slots = vocabulary.getSlots();
          int pieces = Math.max(1, (int) (fileLength / chunkSize));
          ArrayList<byte[]> keys = new ArrayList<byte[]>();
  
          for (int i = 1; i < pieces; ++i) {
              float fraction = (float) i / pieces;
              int slot = (int) (fraction * slots.size());
              keys.add(slots.get(slot).termData);
          }
  
          for (int i = 0; i < pieces; ++i) {
              byte[] firstKey = new byte[0];
              byte[] lastKey = new byte[0];
  
              if (i > 0) {
                  firstKey = keys.get(i - 1);
              }
              if (i < pieces - 1) {
                  lastKey = keys.get(i);
              }
              DocumentSplit split = new DocumentSplit(fileName, fileType, false, firstKey, lastKey);
              processor.process(split);
          }
      }
      
      private void processFile(String fileName) throws IOException {
         // First, try to detect what kind of file this is:
         boolean isCompressed = fileName.endsWith(".gz");
         String fileType = null;
         
         // We'll try to detect by extension first, so we don't have to open the file
         String extension = getExtension(fileName);
         if (extension.equals("corpus") ||
             extension.equals("trecweb") ||
             extension.equals("trectext") ||
             extension.equals("arc") ||
             extension.equals("txt") ||
             extension.equals("html") ||
             extension.equals("xml")) {
             fileType = extension;
         } else {
             // Oh well, we need to autodetect the file type.
             if (IndexReader.isIndexFile(fileName)) {
                 fileType = "corpus";
             } else {
                 // Eventually it'd be nice to do more format detection here.
                 throw new IOException("Couldn't determine file type for: " + fileName);
             }
         }
         
         if (fileType.equals("corpus")) {
             processCorpusFile(fileName, fileType);            
         } else {
             processSplit(fileName, fileType, isCompressed);
         }
     }
     
     private void processDirectory(File root) throws IOException {
         for (File file : root.listFiles()) {
             if (file.isHidden()) {
                 continue;
             }
             if (file.isDirectory()) {
                 processDirectory(file);
             } else {
                 processFile(file.getAbsolutePath());
             }
         }
     }
     
     public void run() throws IOException {
         if (parameters.getXML().containsKey("directory")) {
             List<Value> directories = parameters.getXML().list("directory");
 
             for (Value directory : directories) {
                 File directoryFile = new File(directory.toString());
                 processDirectory(directoryFile);
             }
         } else if (parameters.getXML().containsKey("filename")) {
             List<Value> files = parameters.getXML().list("filename");
 
             for (Value file : files) {
                 processFile(file.toString());
             }
         }
 
         processor.close();
     }
 
     public void setProcessor(Step processor) throws IncompatibleProcessorException {
         Linkage.link(this, processor);
     }
 
     public static void verify(TupleFlowParameters parameters, ErrorHandler handler) {
         FileSource.verify(parameters, handler);
     }
 
     private void processSplit(String fileName, String fileType, boolean isCompressed) throws IOException {
         DocumentSplit split = new DocumentSplit(fileName, fileType, isCompressed, new byte[0], new byte[0]);
         processor.process(split);
     }
 }