/* Copyright (c) 2004 The Nutch Organization.  All rights reserved.   */
/* Use subject to the conditions in http://www.nutch.org/LICENSE.txt. */

package net.nutch.tools;

import java.io.File;
import java.io.FileFilter;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.logging.Logger;

import net.nutch.fetcher.FetcherOutput;
import net.nutch.indexer.DeleteDuplicates;
import net.nutch.indexer.IndexMerger;
import net.nutch.indexer.IndexSegment;
import net.nutch.io.ArrayFile;
import net.nutch.io.Writable;
import net.nutch.parse.ParseData;
import net.nutch.parse.ParseText;
import net.nutch.protocol.Content;
import net.nutch.util.*;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;

/**
 * This class cleans up accumulated segments data, and merges them
 * into a single segment, with no duplicates in it. It uses a "master"
 * unique index of all documents, which either must already exist
 * (by running IndexSegment for each segment, then DeleteDuplicates,
 * and finally IndexMerger), OR the tool can create it just before
 * merging, including creation of per segment sub-indices as needed.
 * <p>The newly created segment is then optionally indexed, so that
 * it can be either merged with more new segments, or used for
 * searching as it is.</p>
 * <p>The original "master" index can be optionally deleted -
 * since it still points to the old segments the new index should
 * be used instead. Old segments may be optionally removed as well,
 * because all needed data has already been copied to the new merged
 * segment.</p>
 * <p>If you use all provided functionality, you can save
 * some manual steps in Nutch operational procedures. After you've
 * run a couple of cycles of fetchlist generation, fetching, DB
 * updating and analyzing, you end up with several segments, possibly
 * containing duplicates. You may then directly run the
 * SegmentMergerTool, with all options turned on, i.e. to first
 * create the master unique index, merge segments into the output
 * segment, index it, and then delete the original segments data and
 * the master index.</p> 
 * 
 * @author Andrzej Bialecki <ab@getopt.org>
 */
public class SegmentMergeTool {
        
  public static final Logger LOG =
    LogFormatter.getLogger("net.nutch.tools.SegmentMergeTool");
  private String master = null;
  private String segments = null;
  private String output = null;
  private File[] segdirs = null;
  private boolean runIndexer = false;
  private boolean createMaster = false;
  private boolean delSegs = false;
  private boolean delMaster = false;
        
  // This class holds together all data readers for an existing segment
  static class SegmentReader {
    public ArrayFile.Reader fetcherReader;
    public ArrayFile.Reader contentReader;
    public ArrayFile.Reader parseTextReader;
    public ArrayFile.Reader parseDataReader;
    public long size = 0L;
                
    public SegmentReader(File dir) throws Exception {
      fetcherReader = new ArrayFile.Reader(new LocalFileSystem(), new File(dir, FetcherOutput.DIR_NAME).toString());
      contentReader = new ArrayFile.Reader(new LocalFileSystem(), new File(dir, Content.DIR_NAME).toString());
      parseTextReader = new ArrayFile.Reader(new LocalFileSystem(), new File(dir, ParseText.DIR_NAME).toString());
      parseDataReader = new ArrayFile.Reader(new LocalFileSystem(), new File(dir, ParseData.DIR_NAME).toString());
      // count the number of valid entries.
      // XXX We assume that all other data files contain the
      // XXX same number of valid entries - which may not be
      // XXX true if Fetcher crashed in the middle of update...
      Writable w = new FetcherOutput();
      while (fetcherReader.next(w) != null)
        size++;
      // reposition to the start
      fetcherReader.reset();
    }
                
    // Close all readers
    public void close() {
      try {
        fetcherReader.close();
      } catch (Exception e) {};
      try {
        contentReader.close();
      } catch (Exception e) {};
      try {
        parseTextReader.close();
      } catch (Exception e) {};
      try {
        parseDataReader.close();
      } catch (Exception e) {};
    }
  }
        
  private HashMap readers = new HashMap();

  // writers for the output segment
  private ArrayFile.Writer fetcherWriter;
  private ArrayFile.Writer contentWriter;
  private ArrayFile.Writer parseTextWriter;
  private ArrayFile.Writer parseDataWriter;
        
  public SegmentMergeTool(String segments, String output, String master,
                          boolean createMaster, boolean runIndexer,
                          boolean delSegs, boolean delMaster) throws Exception {
    this.master = master;
    this.segments = segments;
    this.createMaster = createMaster;
    this.runIndexer = runIndexer;
    this.delSegs = delSegs;
    this.delMaster = delMaster;
    File segs = new File(segments);
    if (!segs.exists() || !segs.isDirectory()) throw new Exception("Not a segments dir: " + segs);
    segdirs = segs.listFiles(new FileFilter() {
        public boolean accept(File file) {
          if (file.isDirectory()) return true;
          return false;
        }
      });
    this.output = output;
  }

  // Create a new segment name
  private String getSegmentName() {
    return new SimpleDateFormat("yyyyMMddHHmmss").format
      (new Date(System.currentTimeMillis()));
  }

  public void run() {
    try {
      // open all segments
      long total = 0L;
      for (int i = 0; i < segdirs.length; i++) {
        SegmentReader sr = new SegmentReader(segdirs[i]);
        total += sr.size;
        LOG.info("Segment " + segdirs[i].getName() + ": "
                 + sr.size + " entries.");
        readers.put(segdirs[i].getName(), sr);
      }
      LOG.info("TOTAL " + total + " input entries in " + segdirs.length + " segments.");
      File masterDir = null;
      if (master != null) masterDir = new File(master);
      LOG.info("Looking for master index in " + masterDir);
      if (masterDir == null || !IndexReader.indexExists(masterDir)) {
        if (!createMaster) {
          LOG.severe("No master index, and createMaster == false");
          return;
        }
        masterDir = new File(new File(segments).getParentFile(), "index");
        if (!masterDir.mkdirs()) {
          LOG.severe("Could not create a master index dir: " + masterDir);
          return;
        }
        LOG.info("Creating master unique index...");
        // check that all segment indexes exist. If not, create them.
        for (int i = 0; i < segdirs.length; i++) {
          File indexerDone = new File(segdirs[i], IndexSegment.DONE_NAME);
          if (!indexerDone.exists()) {
            // Index this segment
            LOG.info(" - creating missing index for " + segdirs[i].getName());
            IndexSegment.main(new String[]{segdirs[i].toString()});
          }
        }
        LOG.info(" - deleting duplicates from indexes in " + segments);
        DeleteDuplicates.main(new String[]{segments, new File(segments).getParentFile().toString()});
        LOG.info(" - creating merged index in " + masterDir);
        String[] args = new String[segdirs.length + 1];
        args[0] = masterDir.toString();
        for (int i = 0; i < segdirs.length; i++) {
          args[i + 1] = segdirs[i].toString();
        }
        IndexMerger.main(args);
      }
      IndexReader ir = IndexReader.open(masterDir);
      File directory = new File(output);
      if (directory.exists() && !directory.isDirectory()) throw new Exception("Output dir is not a directory: " + directory);

      if (!directory.exists()) directory.mkdirs();
      directory = new File(directory, getSegmentName());
      LOG.info("Merging all segments into " + directory);
      directory.mkdirs();
      fetcherWriter = new ArrayFile.Writer
        (new LocalFileSystem(), new File(directory, FetcherOutput.DIR_NAME).toString(),
         FetcherOutput.class);
      contentWriter = new ArrayFile.Writer
        (new LocalFileSystem(), new File(directory, Content.DIR_NAME).toString(), Content.class);
      parseTextWriter = new ArrayFile.Writer
        (new LocalFileSystem(), new File(directory, ParseText.DIR_NAME).toString(), ParseText.class);
      parseDataWriter = new ArrayFile.Writer
        (new LocalFileSystem(), new File(directory, ParseData.DIR_NAME).toString(), ParseData.class);

      FetcherOutput fo = new FetcherOutput();
      Content co = new Content();
      ParseText pt = new ParseText();
      ParseData pd = new ParseData();

      int docCnt = ir.numDocs();
      int outputCnt = 0;
      for (int i = 0; i < docCnt; i++) {
        if (i > 0 && (i % 500 == 0)) LOG.info("Processed " + i + " entries.");
        if (ir.isDeleted(i)) {
          LOG.fine("\n- skip deleted doc # " + i);
          continue;
        }
        Document doc = ir.document(i);
        String segName = doc.get("segment");
        SegmentReader sr = (SegmentReader)readers.get(segName);
        if (sr == null) {
          LOG.warning("\n- no SegmentReader for " + segName);
          continue;
        }
        String docNo = doc.get("docNo");
        long docid = -1L;
        try {
          docid = Long.parseLong(docNo, 16);
        } catch (Exception e) {
          LOG.warning("\n- bad docNo: " + docNo);
          continue;
        }
        // get data from the reader
        sr.fetcherReader.get(docid, fo);
        sr.contentReader.get(docid, co);
        sr.parseTextReader.get(docid, pt);
        sr.parseDataReader.get(docid, pd);
        // write it back
        fetcherWriter.append(fo);
        contentWriter.append(co);
        parseTextWriter.append(pt);
        parseDataWriter.append(pd);
        outputCnt++;
      }
      ir.close();
      fetcherWriter.close();
      contentWriter.close();
      parseTextWriter.close();
      parseDataWriter.close();
      for (Iterator i = readers.keySet().iterator(); i.hasNext(); ) {
        SegmentReader sr = (SegmentReader)readers.get(i.next());
        sr.close();
      }
      if (runIndexer) {
        LOG.info("Creating new segment index...");
        IndexSegment.main(new String[]{directory.toString()});
      }
      if (delSegs) {
        LOG.info("Deleting old segments...");
        for (int i = 0; i < segdirs.length; i++) {
          FileUtil.fullyDelete(segdirs[i]);
        }
      }
      if (delMaster) {
        LOG.info("Deleting old master index...");
        FileUtil.fullyDelete(masterDir);
      }
      LOG.info("DONE segment merging, INPUT: " + total + " -> OUTPUT: " + outputCnt + " entries.");
    } catch (Exception e) {
      e.printStackTrace();
      LOG.severe(e.getMessage());
    }
  }
        
  public static void main(String[] args) throws Exception {
    if (args.length < 2) {
      System.err.println("Too few arguments.\n");
      usage();
      System.exit(-1);
    }
    boolean createMaster = false;
    boolean runIndexer = false;
    boolean delSegs = false;
    boolean delMaster = false;
    String index = null;
    String output = null;
    for (int i = 1; i < args.length; i++) {
      if (args[i].equals("-m")) {
        if (args.length > i + 1) {
          index = args[++i];
          continue;
        } else {
          System.err.println("Required value of '-m' argument missing.\n");
          usage();
          System.exit(-1);
        }
      } else if(args[i].equals("-o")) {
        if (args.length > i + 1) {
          output = args[++i];
          continue;
        } else {
          System.err.println("Required value of '-o' argument missing.\n");
          usage();
          System.exit(-1);
        }
      }
      else if (args[i].equals("-cm")) createMaster = true;
      else if (args[i].equals("-i")) runIndexer = true;
      else if (args[i].equals("-ds")) delSegs = true;
      else if (args[i].equals("-dm")) delMaster = true;
                        
    }
    if (index == null && !createMaster) {
      System.err.println("No master index, and createMaster == false.\n");
      usage();
      System.exit(-1);
    }
    if (output == null) output = args[0];
    SegmentMergeTool st = new SegmentMergeTool(args[0], output, index,
                                               createMaster, runIndexer, delSegs, delMaster);
    st.run();
  }
        
  private static void usage() {
    System.err.println("SegmentMergeTool <input_segments_dir> [-o <output_segment_dir>] [-m <unique_index_dir> | -cm] [-i] [-ds] [-dm]");
    System.err.println("\t<input_segments_dir>\tpath to directory containing\n\t\t\t\tall input segments");
    System.err.println("\t-o <output_segment_dir>\t(optional) path to directory which will\n\t\t\t\tcontain a single output segment.\n\t\t\tNOTE: If not present, the original segments path will be used.");
    System.err.println("\n\t-m <unique_index_dir>\tpath to 'master' unique\n\t\t\t\tindex for all input segments.\n\t\t\tNOTE: either this or the '-cm' option MUST be used!");
    System.err.println("\t-cm\t\t(optional) create the 'master' index first,\n\t\t\tby running IndexSegment on each segment,\n\t\t\tthen DeleteDuplicates, and then IndexMerger.");
    System.err.println("\t-i\t\t(optional) index the output segment.");
    System.err.println("\t-ds\t\t(optional) delete the original segments when finished.");
    System.err.println("\t-dm\t\t(optional) delete the original master index when finished.");
    System.err.println();
  }
}