package org.infinispan.persistence.leveldb;

import org.infinispan.commons.CacheConfigurationException;
import org.infinispan.commons.util.Util;
import org.infinispan.persistence.leveldb.configuration.LevelDBStoreConfiguration;
import org.infinispan.persistence.leveldb.logging.Log;
import org.infinispan.marshall.core.MarshalledEntry;
import org.infinispan.metadata.InternalMetadata;
import org.infinispan.persistence.CacheLoaderException;
import org.infinispan.persistence.PersistenceUtil;
import org.infinispan.persistence.TaskContextImpl;
import org.infinispan.persistence.spi.AdvancedLoadWriteStore;
import org.infinispan.persistence.spi.InitializationContext;
import org.infinispan.util.logging.LogFactory;
import org.iq80.leveldb.CompressionType;
import org.iq80.leveldb.DB;
import org.iq80.leveldb.DBException;
import org.iq80.leveldb.DBFactory;
import org.iq80.leveldb.DBIterator;
import org.iq80.leveldb.Options;
import org.iq80.leveldb.ReadOptions;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.LinkedBlockingQueue;

public class LevelDBStore implements AdvancedLoadWriteStore {
   private static final Log log = LogFactory.getLog(LevelDBStore.class, Log.class);

   private static final String JNI_DB_FACTORY_CLASS_NAME = "org.fusesource.leveldbjni.JniDBFactory";
   private static final String JAVA_DB_FACTORY_CLASS_NAME = "org.iq80.leveldb.impl.Iq80DBFactory";
   private static final String[] DB_FACTORY_CLASS_NAMES = new String[] { JNI_DB_FACTORY_CLASS_NAME, JAVA_DB_FACTORY_CLASS_NAME };

   private LevelDBStoreConfiguration configuration;
   private BlockingQueue<ExpiryEntry> expiryEntryQueue;
   private DBFactory dbFactory;
   private DB db;
   private DB expiredDb;
   private InitializationContext ctx;


   @Override
   public void init(InitializationContext ctx) {
      this.configuration = ctx.getConfiguration();
      this.dbFactory = newDbFactory();
      this.ctx = ctx;

      if (this.dbFactory == null) {
         throw log.cannotLoadlevelDBFactories(Arrays.toString(DB_FACTORY_CLASS_NAMES));
      }
      String dbFactoryClassName = this.dbFactory.getClass().getName();
      if (dbFactoryClassName.equals("org.iq80.leveldb.impl.Iq80DBFactory")) {
         log.infoUsingJavaDbFactory(dbFactoryClassName);
      } else {
         log.infoUsingJNIDbFactory(dbFactoryClassName);
      }
   }

   protected DBFactory newDbFactory() {
      switch (configuration.implementationType()) {
         case JNI: {
            return Util.getInstance(JNI_DB_FACTORY_CLASS_NAME, LevelDBStore.class.getClassLoader());
         }
         case JAVA: {
            return Util.getInstance(JAVA_DB_FACTORY_CLASS_NAME, LevelDBStore.class.getClassLoader());
         }
         default: {
            for (String className : DB_FACTORY_CLASS_NAMES) {
               try {
                  return Util.getInstance(className, LevelDBStore.class.getClassLoader());
               } catch (Throwable e) {
                  if (log.isDebugEnabled())
                     log.debugUnableToInstantiateDbFactory(className, e);
               }
            }
         }
      }
      return null;
   }

   @Override
   public void start()  {
      expiryEntryQueue = new LinkedBlockingQueue<ExpiryEntry>(configuration.expiryQueueSize());

      try {
         db = openDatabase(getQualifiedLocation(), dataDbOptions());
         expiredDb = openDatabase(getQualifiedExpiredLocation(), expiredDbOptions());
      } catch (IOException e) {
         throw new CacheConfigurationException("Unable to open database", e);
      }
   }

   private String sanitizedCacheName() {
      String cacheFileName = ctx.getCache().getName().replaceAll("[^a-zA-Z0-9-_\\.]", "_");
      return cacheFileName;
   }

   private String getQualifiedLocation() {
      return configuration.location() + sanitizedCacheName();
   }

   private String getQualifiedExpiredLocation() {
      return configuration.expiredLocation() + sanitizedCacheName();
   }

   private Options dataDbOptions() {
      Options options = new Options().createIfMissing(true);

      options.compressionType(CompressionType.valueOf(configuration.compressionType().name()));

      if (configuration.blockSize() != null) {
         options.blockSize(configuration.blockSize());
      }

      if (configuration.cacheSize() != null) {
         options.cacheSize(configuration.cacheSize());
      }

      return options;
   }

   private Options expiredDbOptions() {
      return new Options().createIfMissing(true);
   }

   /**
    * Creates database if it doesn't exist.
    */
   protected DB openDatabase(String location, Options options) throws IOException {
      File dir = new File(location);

      // LevelDB JNI Option createIfMissing doesn't seem to work properly
      dir.mkdirs();
      return dbFactory.open(dir, options);
   }

   protected void destroyDatabase(String location) throws IOException {
      File dir = new File(location);
      dbFactory.destroy(dir, new Options());
   }

   protected DB reinitDatabase(String location, Options options) throws IOException {
      destroyDatabase(location);
      return openDatabase(location, options);
   }

   protected void reinitAllDatabases() throws IOException {
      // assuming this is only called from within clearLockSafe()
      // clear() acquires a global lock before calling clearLockSafe()
      try {
         db.close();
      } catch (IOException e) {
         log.warnUnableToCloseDb(e);
      }
      try {
         expiredDb.close();
      } catch (IOException e) {
         log.warnUnableToCloseExpiredDb(e);
      }
      db = reinitDatabase(getQualifiedLocation(), dataDbOptions());
      expiredDb = reinitDatabase(getQualifiedExpiredLocation(), expiredDbOptions());
   }

   @Override
   public void stop()  {
      try {
         db.close();
      } catch (IOException e) {
         log.warnUnableToCloseDb(e);
      }

      try {
         expiredDb.close();
      } catch (IOException e) {
         log.warnUnableToCloseExpiredDb(e);
      }
   }

   @Override
   public void clear() {
      long count = 0;
      DBIterator it = db.iterator(new ReadOptions().fillCache(false));
      boolean destroyDatabase = false;

      if (configuration.clearThreshold() <= 0) {
         try {
            for (it.seekToFirst(); it.hasNext();) {
               Map.Entry<byte[], byte[]> entry = it.next();
               db.delete(entry.getKey());
               count++;

               if (count > configuration.clearThreshold()) {
                  destroyDatabase = true;
                  break;
               }
            }
         } finally {
            try {
               it.close();
            } catch (IOException e) {
               log.warnUnableToCloseDbIterator(e);
            }
         }
      } else {
         destroyDatabase = true;
      }

      if (destroyDatabase) {
         try {
            reinitAllDatabases();
         } catch (IOException e) {
            throw new CacheLoaderException(e);
         }
      }
   }

   @Override
   public int size() {
      return PersistenceUtil.count(this, null);
   }

   @Override
   public boolean contains(Object key) {
      try {
         return load(key) != null;
      } catch (Exception e) {
         throw new CacheLoaderException(e);
      }
   }

   @SuppressWarnings("unchecked")
   @Override
   public void process(KeyFilter keyFilter, CacheLoaderTask cacheLoaderTask, Executor executor, boolean loadValues, boolean loadMetadata) {

      int batchSize = 100;
      ExecutorCompletionService ecs = new ExecutorCompletionService(executor);
      int tasks = 0;
      final TaskContext taskContext = new TaskContextImpl();

      List<Map.Entry<byte[], byte[]>> entries = new ArrayList<Map.Entry<byte[], byte[]>>(batchSize);
      DBIterator it = db.iterator(new ReadOptions().fillCache(false));
      try {
         for (it.seekToFirst(); it.hasNext();) {
            Map.Entry<byte[], byte[]> entry = it.next();
            entries.add(entry);
            if (entries.size() == batchSize) {
               final List<Map.Entry<byte[], byte[]>> batch = entries;
               entries = new ArrayList<Map.Entry<byte[], byte[]>>(batchSize);
               submitProcessTask(cacheLoaderTask, keyFilter,ecs, taskContext, batch);
               tasks++;
            }
         }
         if (!entries.isEmpty()) {
            submitProcessTask(cacheLoaderTask, keyFilter,ecs, taskContext, entries);
            tasks++;
         }

         PersistenceUtil.waitForAllTasksToComplete(ecs, tasks);
      } catch (Exception e) {
         throw new CacheLoaderException(e);
      } finally {
         try {
            it.close();
         } catch (IOException e) {
            log.warnUnableToCloseDbIterator(e);
         }
      }
   }

   @SuppressWarnings("unchecked")
   private void submitProcessTask(final CacheLoaderTask cacheLoaderTask, final KeyFilter filter, ExecutorCompletionService ecs,
                                  final TaskContext taskContext, final List<Map.Entry<byte[], byte[]>> batch) {
      ecs.submit(new Callable<Void>() {
         @Override
         public Void call() throws Exception {
            long now = ctx.getTimeService().wallClockTime();
            for (Map.Entry<byte[], byte[]> entry : batch) {
               if (taskContext.isStopped()) {break;}
               Object key = unmarshall(entry.getKey());
               if (filter == null || filter.shouldLoadKey(key)) {
                  MarshalledEntry unmarshall = (MarshalledEntry) unmarshall(entry.getValue());
                  boolean isExpired = unmarshall.getMetadata() != null && unmarshall.getMetadata().isExpired(now);
                  if (!isExpired) {
                     cacheLoaderTask.processEntry(unmarshall, taskContext);
                  }
               }
            }
            return null;
         }
      });
   }

   @Override
   public boolean delete(Object key)  {
      try {
         byte[] keyBytes = marshall(key);
         if (db.get(keyBytes) == null) {
            return false;
         }
         db.delete(keyBytes);
         return true;
      } catch (Exception e) {
         throw new CacheLoaderException(e);
      }
   }

   @Override
   public void write(MarshalledEntry me)  {
      try {
         db.put(marshall(me.getKey()), marshall(me));
         InternalMetadata meta = me.getMetadata();
         if (meta != null && meta.expiryTime() > -1) {
            addNewExpiry(me);
         }
      } catch (Exception e) {
         throw new DBException(e);
      }
   }

   @Override
   public MarshalledEntry load(Object key)  {
      try {
         MarshalledEntry me = (MarshalledEntry) unmarshall(db.get(marshall(key)));
         if (me == null) return null;

         InternalMetadata meta = me.getMetadata();
         if (meta != null && meta.isExpired(ctx.getTimeService().wallClockTime())) {
            return null;
         }
         return me;
      } catch (Exception e) {
         throw new CacheLoaderException(e);
      }
   }

   @SuppressWarnings("unchecked")
   @Override
   public void purge(Executor executor, PurgeListener purgeListener) {
      try {
         // Drain queue and update expiry tree
         List<ExpiryEntry> entries = new ArrayList<ExpiryEntry>();
         expiryEntryQueue.drainTo(entries);
         for (ExpiryEntry entry : entries) {
            final byte[] expiryBytes = marshall(entry.expiry);
            final byte[] keyBytes = marshall(entry.key);
            final byte[] existingBytes = expiredDb.get(expiryBytes);

            if (existingBytes != null) {
               // in the case of collision make the key a List ...
               final Object existing = unmarshall(existingBytes);
               if (existing instanceof List) {
                  ((List<Object>) existing).add(entry.key);
                  expiredDb.put(expiryBytes, marshall(existing));
               } else {
                  List<Object> al = new ArrayList<Object>(2);
                  al.add(existing);
                  al.add(entry.key);
                  expiredDb.put(expiryBytes, marshall(al));
               }
            } else {
               expiredDb.put(expiryBytes, keyBytes);
            }
         }

         List<Long> times = new ArrayList<Long>();
         List<Object> keys = new ArrayList<Object>();
         DBIterator it = expiredDb.iterator(new ReadOptions().fillCache(false));
         try {
            for (it.seekToFirst(); it.hasNext();) {
               Map.Entry<byte[], byte[]> entry = it.next();

               Long time = (Long) unmarshall(entry.getKey());
               if (time > System.currentTimeMillis())
                  break;
               times.add(time);
               Object key = unmarshall(entry.getValue());
               if (key instanceof List)
                  keys.addAll((List<?>) key);
               else
                  keys.add(key);
            }

            for (Long time : times) {
               expiredDb.delete(marshall(time));
            }

            if (!keys.isEmpty())
               log.debugf("purge (up to) %d entries", keys.size());
            int count = 0;
            long currentTimeMillis = System.currentTimeMillis();
            for (Object key : keys) {
               byte[] keyBytes = marshall(key);

               byte[] b = db.get(keyBytes);
               if (b == null)
                  continue;
               MarshalledEntry me = (MarshalledEntry) ctx.getMarshaller().objectFromByteBuffer(b);
               if (me.getMetadata() != null && me.getMetadata().isExpired(ctx.getTimeService().wallClockTime())) {
                  // somewhat inefficient to FIND then REMOVE...
                  db.delete(keyBytes);
                  count++;
               }
            }
            if (count != 0)
               log.debugf("purged %d entries", count);
         } catch (Exception e) {
            throw new CacheLoaderException(e);
         } finally {
            try {
               it.close();
            } catch (IOException e) {
               log.warnUnableToCloseDbIterator(e);
            }
         }
      } catch (CacheLoaderException e) {
         throw e;
      } catch (Exception e) {
         throw new CacheLoaderException(e);
      }
   }

   private byte[] marshall(Object entry) throws IOException, InterruptedException {
      return ctx.getMarshaller().objectToByteBuffer(entry);
   }

   private Object unmarshall(byte[] bytes) throws IOException, ClassNotFoundException {
      if (bytes == null)
         return null;

      return ctx.getMarshaller().objectFromByteBuffer(bytes);
   }

   private void addNewExpiry(MarshalledEntry entry) throws IOException {
      long expiry = entry.getMetadata().expiryTime();
      long maxIdle = entry.getMetadata().maxIdle();
      if (maxIdle > 0) {
         // Coding getExpiryTime() for transient entries has the risk of
         // being a moving target
         // which could lead to unexpected results, hence, InternalCacheEntry
         // calls are required
         expiry = maxIdle + System.currentTimeMillis();
      }
      Long at = expiry;
      Object key = entry.getKey();

      try {
         expiryEntryQueue.put(new ExpiryEntry(at, key));
      } catch (InterruptedException e) {
         Thread.currentThread().interrupt(); // Restore interruption status
      }
   }

   private static final class ExpiryEntry {
      private final Long expiry;
      private final Object key;

      private ExpiryEntry(long expiry, Object key) {
         this.expiry = expiry;
         this.key = key;
      }

      @Override
      public int hashCode() {
         final int prime = 31;
         int result = 1;
         result = prime * result + ((key == null) ? 0 : key.hashCode());
         return result;
      }

      @Override
      public boolean equals(Object obj) {
         if (this == obj)
            return true;
         if (obj == null)
            return false;
         if (getClass() != obj.getClass())
            return false;
         ExpiryEntry other = (ExpiryEntry) obj;
         if (key == null) {
            if (other.key != null)
               return false;
         } else if (!key.equals(other.key))
            return false;
         return true;
      }

   }

}
