package com.atlassian.bonnie.index;
   
   import org.apache.log4j.Category;
   import org.apache.lucene.analysis.Analyzer;
   import org.apache.lucene.document.Document;
   import org.apache.lucene.index.IndexWriter;
   import org.apache.lucene.store.Directory;
   import org.apache.lucene.store.FSDirectory;
   
  import java.io.File;
  import java.io.IOException;
  import java.util.*;
  
  /**
   * Writes documents to multiple temporary indices.
   */
  public class TempIndexWriter
  {
      protected static Category log = Category.getInstance(TempIndexWriter.class);
  
      private final Map writers = Collections.synchronizedMap(new HashMap());
      private Analyzer analyzerForIndexing;
      private String tmpDir;
  
      public TempIndexWriter(Analyzer analyzerForIndexing, String tmpDir)
      {
          this.analyzerForIndexing = analyzerForIndexing;
          this.tmpDir = tmpDir;
      }
  
      /**
       * Add a document to an index.
       * @param key key representing which index to write to. A new writer will be created if no writer exists for the key.
       * @param doc document to add
       * @throws IOException
       */
      public void addDocument(String key, Document doc) throws IOException
      {
          WriterData writerData = (WriterData) writers.get(key);
          if (writerData == null)
          {
              writerData = createData();
              writers.put(key, writerData);
          }
          writerData.writer.addDocument(doc);
      }
  
      /**
       * Delete the respective temp indices.
       * @param prefix
       * @throws IOException
       */
      public void close(String prefix) throws IOException
      {
          for (Iterator it = writers.keySet().iterator(); it.hasNext();)
          {
              String key = it.next().toString();
              if (key.startsWith(prefix))
              {
                  WriterData writerData = (WriterData) writers.get(key);
                  if (writerData != null) {
                      delete(writerData.tmpIndexDir);
                  }
              }
          }
      }
  
      /**
       * Delete the respective temp indices.
       * @throws IOException
       */
      public void closeAll() throws IOException
      {
          for (Iterator it = writers.keySet().iterator(); it.hasNext();)
          {
              WriterData writerData = (WriterData) writers.get(it.next());
              delete(writerData.tmpIndexDir);
          }
          writers.clear();
      }
  
      /**
       * Close all temporary writers, and merge all temporary indices to the writer provided.
       * @param writer
       * @throws IOException
       */
      public void merge(IndexWriter writer) throws IOException
      {
          final List indexDirectories = new LinkedList();
          for (Iterator it = writers.keySet().iterator(); it.hasNext();)
          {
              WriterData writerData = (WriterData) writers.get(it.next());
              writerData.writer.close();
              indexDirectories.add(writerData.dir);
          }
          writer.addIndexes((Directory[]) indexDirectories.toArray(new Directory[indexDirectories.size()]));
      }
  
      /**
      * Close relevant temporary writers, and merge temporary indices identified by keys which start with the provided prefix to the
      * writer provided.
      * @param prefix key prefix
      * @param writer
      * @throws IOException
      */
     public void merge(String prefix, IndexWriter writer) throws IOException
     {
         final List indexDirectories = new LinkedList();
         for (Iterator it = writers.keySet().iterator(); it.hasNext();)
         {
             String key = it.next().toString();
             if (key.startsWith(prefix))
             {
                 WriterData writerData = (WriterData) writers.get(key);
                 writerData.writer.close();
                 indexDirectories.add(writerData.dir);
             }
         }
         System.out.println("Index directories:" + indexDirectories);
         if (indexDirectories.size() > 0)
             writer.addIndexes((Directory[]) indexDirectories.toArray(new Directory[indexDirectories.size()]));
     }
 
     protected final boolean delete(File directory)
     {
         File[] files = directory.listFiles();
         for (int i = 0; i < files.length; i++)
         {
             if (!files[i].delete())
             {
                 log.error("Failed to delete index file: " + files[i].getAbsolutePath());
                 return false;
             }
         }
         return directory.delete();
     }
 
     private WriterData createData() throws IOException
     {
         // create the temp index directory.
         File tmpIndexDir = getTmpDir() == null ? File.createTempFile("lucene", "index") :
                 File.createTempFile("lucene", "index", new File(getTmpDir()));
         if (!tmpIndexDir.delete() || !tmpIndexDir.mkdirs())
         {
             throw new IOException("Unable to create temporary index directory: " + tmpIndexDir);
         }
         Directory dir = FSDirectory.getDirectory(tmpIndexDir);
         //TODO: This probably doesn't need to be here. We should just construct the next IndexWriter with create=true
         new IndexWriter(dir, null, true).close();
         IndexWriter writer = new IndexWriter(dir, analyzerForIndexing, false);
         writer.setMergeFactor(50);
         writer.setMaxBufferedDocs(300);
         writer.setMaxMergeDocs(Integer.MAX_VALUE);    // better for fast searches and batch indexing, but not if we need small segments
 
         return new WriterData(writer, dir, tmpIndexDir);
     }
 
     public String getTmpDir()
     {
         return tmpDir;
     }
 
     protected class WriterData
     {
         IndexWriter writer;
         Directory dir;
         File tmpIndexDir;
 
         protected WriterData(IndexWriter writer, Directory dir, File tmpIndexDir)
         {
             this.writer = writer;
             this.dir = dir;
             this.tmpIndexDir = tmpIndexDir;
         }
     }
 }