/**
 * OrientDB client binding for YCSB.
 *
 * Submitted by Luca Garulli on 5/10/2012.
 *
 */

package com.yahoo.ycsb.db;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;
import java.util.Random;
import java.util.List;
import java.util.ArrayList;
import java.lang.Integer;

import com.orientechnologies.orient.client.remote.OServerAdmin;
import com.orientechnologies.orient.core.config.OGlobalConfiguration;
import com.orientechnologies.orient.core.db.document.ODatabaseDocumentTx;
import com.orientechnologies.orient.core.db.document.ODatabaseDocumentPool;
import com.orientechnologies.orient.core.db.ODatabaseComplex.OPERATION_MODE;
import com.orientechnologies.orient.core.db.ODatabaseRecordThreadLocal;
import com.orientechnologies.orient.core.storage.*;
import com.orientechnologies.orient.core.id.OClusterPosition;
import com.orientechnologies.orient.core.id.ORecordId;
import com.orientechnologies.orient.core.dictionary.ODictionary;
import com.orientechnologies.orient.core.intent.OIntentMassiveInsert;
import com.orientechnologies.orient.core.record.ORecordInternal;
import com.orientechnologies.orient.core.record.impl.ODocument;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.StringByteIterator;
import com.yahoo.ycsb.Host;

/**
 * OrientDB client for YCSB framework.
 * 
 * Properties to set:
 * 
 * orientdb.url=local:C:/temp/databases/ycsb or remote:localhost:2424 <br>
 * orientdb.database=ycsb <br>
 * orientdb.user=admin <br>
 * orientdb.password=admin <br>
 * 
 * @author Luca Garulli
 * 
 */
public class OrientDBClient extends DB {

  // private ODatabaseDocumentTx              db;
  private static final String              CLASS = "usertable";
  private ODictionary<ORecordInternal<?>>  dictionary;
  private OServerAdmin                     serverAdmin;
  private List<Host>                       listHosts;
  
  private Map<String, ODatabaseDocumentTx> databases;
  private Map<String, ODictionary<ORecordInternal<?>>> dictionaries;

  private String user;
  private String password;
  private Boolean newdb;

  public static final int MIN_KEY = 0;
  public static final int MAX_KEY = 10;

  /**
   * Initialize any state for this DB. Called once per DB instance; there is one DB instance per client thread.
   */
  public void init() throws DBException {
    // initialize OrientDB driver
    Properties props = getProperties();

    String urls = props.getProperty("orientdb.urls", "remote:localhost/ycsb");
    user = props.getProperty("orientdb.user", "rchow");
    password = props.getProperty("orientdb.password", "hiraymond");
    newdb = Boolean.parseBoolean(props.getProperty("orientdb.newdb", "false"));

    String[] all_urls = urls.split(",");

    // Set bucket interval
    int interval = (MAX_KEY - MIN_KEY) / all_urls.length;

    // List to store hosts
    listHosts = new ArrayList<Host>();

    databases = new HashMap<String, ODatabaseDocumentTx>(all_urls.length);
    dictionaries = new HashMap<String, ODictionary<ORecordInternal<?>>>(all_urls.length);

    // Populate hosts and databases
    for( int i = 0; i < all_urls.length; i++ )
    {
      listHosts.add( new Host(i*interval, (i+1)*interval, all_urls[i]) );
      databases.put( all_urls[i], new ODatabaseDocumentTx(all_urls[i]) );
    }

    try {

      // System.out.println("OrientDB loading database url = " + url);

      OGlobalConfiguration.STORAGE_KEEP_OPEN.setValue(false);
      // OGlobalConfiguration.CLIENT_CHANNEL_MAX_POOL.setValue(5000);
      // OGlobalConfiguration.NETWORK_LOCK_TIMEOUT.setValue(10000);
      
      // Set up the DB
      setupDB();

      // System.out.println("OrientDB connection created with " + url);

    } catch (Exception e1) {
      System.err.println("Could not initialize OrientDB connection pool for Loader: " + e1.toString());
      e1.printStackTrace();
      return;
    }

  }

  public void setupDB()
  {

    try {

      for( Entry<String, ODatabaseDocumentTx> entry : databases.entrySet() )
      {
        String url = entry.getKey();
        ODatabaseDocumentTx database = entry.getValue();

        // Establish a remote connection to the host (url)
        serverAdmin = new OServerAdmin(url).connect(user, password);

        // Check if database exists
        if( serverAdmin.existsDatabase() )
        {
          // Re-create DB
          if(newdb)
          {
            // System.out.println("OrientDB drop and recreate fresh db");
            serverAdmin.connect(user, password).dropDatabase().close();
            serverAdmin.connect(user, password).createDatabase("document", "local").close();
          }
        }
        // Database does not exist, create new one
        else
        {
          // System.out.println("OrientDB database not found, create fresh db");
          serverAdmin.connect(user, password).createDatabase("document", "local").close();
        }

        // Set DB properties
        database.setProperty("minPool", 5);
        database.setProperty("maxPool", 200);
        
        // Open DB
        database.open(user, password);

        // Get dictionary for DB
        ODictionary<ORecordInternal<?>> dt = database.getMetadata().getIndexManager().getDictionary();

        // Add dictionary
        dictionaries.put( url, dt );

        // Check if CLASS exists else create
        if (!database.getMetadata().getSchema().existsClass(CLASS))
          database.getMetadata().getSchema().createClass(CLASS);

        // Tune DB for massive insertions
        database.declareIntent(new OIntentMassiveInsert());

        entry.setValue(database);
      } 

    } catch (Exception e1) {
      System.err.println("Could not initialize OrientDB connection pool for Loader: " + e1.toString());
      e1.printStackTrace();
      return;
    }
  }

  public String findHost(String key)
  {
    // Find first number of key
    String strIntKey = key.substring(4, 5);

    // Convert string to integer
    int intKey = Integer.parseInt(strIntKey);

    // Find which host key belongs to
    Host host = Host.whichHost( intKey, listHosts );

    return host.getHostIP();
  }

  @Override
  public void cleanup() throws DBException {

    for( Entry<String, ODatabaseDocumentTx> entry : databases.entrySet() )
    {
      ODatabaseDocumentTx database = entry.getValue();
      
      if( database != null )
      {
        database.close();
        database = null;
        entry.setValue(database);
      }
    }
  }

  @Override
  /**
   * Insert a record in the database. Any field/value pairs in the specified values HashMap will be written into the record with the specified
   * record key.
   *
   * @param table The name of the table
   * @param key The record key of the record to insert.
   * @param values A HashMap of field/value pairs to insert in the record
   * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
   */
  public int insert(String table, String key, HashMap<String, ByteIterator> values) {

    String hostIP = findHost(key);

    try {

      ODatabaseDocumentTx database = databases.get(hostIP);

      ODatabaseRecordThreadLocal.INSTANCE.set(database);

      final ODocument document = new ODocument(CLASS);
      for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet())
        document.field(entry.getKey(), entry.getValue());
      document.save();
      dictionaries.get(hostIP).put(key, document);

      return 0;
    } catch (Exception e) {
      e.printStackTrace();
    }

    return 1;
  }

  @Override
  /**
   * Delete a record from the database.
   *
   * @param table The name of the table
   * @param key The record key of the record to delete.
   * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
   */
  public int delete(String table, String key) {
    try {

      String hostIP = findHost(key);

      dictionaries.get(hostIP).remove(key);

      return 0;
    } catch (Exception e) {
      e.printStackTrace();
    }
    return 1;
  }

  @Override
  /**
   * Read a record from the database. Each field/value pair from the result will be stored in a HashMap.
   *
   * @param table The name of the table
   * @param key The record key of the record to read.
   * @param fields The list of fields to read, or null for all of them
   * @param result A HashMap of field/value pairs for the result
   * @return Zero on success, a non-zero error code on error or "not found".
   */
  public int read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
    try {
      
      String hostIP = findHost(key);

      final ODocument document = dictionaries.get(hostIP).get(key);
      if (document != null) {
        if (fields != null)
          for (String field : fields)
            result.put(field, new StringByteIterator((String) document.field(field)));
          else
            for (String field : document.fieldNames())
              result.put(field, new StringByteIterator((String) document.field(field)));
            return 0;
          }
        } catch (Exception e) {
          e.printStackTrace();
        }
        return 1;
      }

      @Override
  /**
   * Update a record in the database. Any field/value pairs in the specified values HashMap will be written into the record with the specified
   * record key, overwriting any existing values with the same field name.
   *
   * @param table The name of the table
   * @param key The record key of the record to write.
   * @param values A HashMap of field/value pairs to update in the record
   * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
   */
  public int update(String table, String key, HashMap<String, ByteIterator> values) {
    try {
      
      String hostIP = findHost(key);

      ODatabaseDocumentTx database = databases.get(hostIP);

      ODatabaseRecordThreadLocal.INSTANCE.set(database);

      final ODocument document = dictionaries.get(hostIP).get(key);
      if (document != null) {
        for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet())
          document.field(entry.getKey(), entry.getValue());
        document.save();
        return 0;
      }
    } catch (Exception e) {
      e.printStackTrace();
    }
    return 1;
  }

  @Override
  /**
   * Perform a range scan for a set of records in the database. Each field/value pair from the result will be stored in a HashMap.
   *
   * @param table The name of the table
   * @param startkey The record key of the first record to read.
   * @param recordcount The number of records to read
   * @param fields The list of fields to read, or null for all of them
   * @param result A Vector of HashMaps, where each HashMap is a set field/value pairs for one record
   * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
   */
  public int scan(String table, String startkey, int recordcount, Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
    try {
      final Collection<ODocument> documents = dictionary.getIndex().getEntriesMajor(startkey, true, recordcount);
      for (ODocument document : documents) {
        final HashMap<String, ByteIterator> entry = new HashMap<String, ByteIterator>(fields.size());
        result.add(entry);

        for (String field : fields)
          entry.put(field, new StringByteIterator((String) document.field(field)));
      }

      return 0;
    } catch (Exception e) {
      e.printStackTrace();
    }
    return 1;
  }
}

