/**
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License. You
 * may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License. See accompanying
 * LICENSE file.
 */

package com.yahoo.ycsb.db;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

import com.yahoo.ycsb.ByteArrayByteIterator;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.Status;
import com.yahoo.ycsb.measurements.Measurements;

import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.BufferedMutator;
import org.apache.hadoop.hbase.client.BufferedMutatorParams;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Delete;
import org.apache.hadoop.hbase.client.Durability;
import org.apache.hadoop.hbase.client.Get;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.ResultScanner;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.filter.PageFilter;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.ConcurrentModificationException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.Vector;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * HBase 1.0 client for YCSB framework.
 *
 * A modified version of HBaseClient (which targets HBase v0.9) utilizing the
 * HBase 1.0.0 API.
 *
 * This client also adds toggleable client-side buffering and configurable write
 * durability.
 */
public class HBaseClient10 extends com.yahoo.ycsb.DB {
  private Configuration config = HBaseConfiguration.create();

  private static AtomicInteger threadCount = new AtomicInteger(0);

  private boolean debug = false;

  private String tableName = "";

  /**
   * A Cluster Connection instance that is shared by all running ycsb threads.
   * Needs to be initialized late so we pick up command-line configs if any.
   * To ensure one instance only in a multi-threaded context, guard access
   * with a 'lock' object.
   * @See #CONNECTION_LOCK.
   */
  private static Connection connection = null;
  private static final Object CONNECTION_LOCK = new Object();

  // Depending on the value of clientSideBuffering, either bufferedMutator
  // (clientSideBuffering) or currentTable (!clientSideBuffering) will be used.
  private Table currentTable = null;
  private BufferedMutator bufferedMutator = null;

  private String columnFamily = "";
  private byte[] columnFamilyBytes;

  /**
   * Durability to use for puts and deletes.
   */
  private Durability durability = Durability.USE_DEFAULT;

  /** Whether or not a page filter should be used to limit scan length. */
  private boolean usePageFilter = true;

  /**
   * If true, buffer mutations on the client. This is the default behavior for
   * HBaseClient. For measuring insert/update/delete latencies, client side
   * buffering should be disabled.
   */
  private boolean clientSideBuffering = false;
  private long writeBufferSize = 1024 * 1024 * 12;

  /**
   * Initialize any state for this DB. Called once per DB instance; there is one
   * DB instance per client thread.
   */
  @Override
  public void init() throws DBException {
    if ("true"
        .equals(getProperties().getProperty("clientbuffering", "false"))) {
      this.clientSideBuffering = true;
    }
    if (getProperties().containsKey("writebuffersize")) {
      writeBufferSize =
          Long.parseLong(getProperties().getProperty("writebuffersize"));
    }

    if (getProperties().getProperty("durability") != null) {
      this.durability =
          Durability.valueOf(getProperties().getProperty("durability"));
    }

    if ("kerberos".equalsIgnoreCase(config.get("hbase.security.authentication"))) {
      config.set("hadoop.security.authentication", "Kerberos");
      UserGroupInformation.setConfiguration(config);
    }

    if ((getProperties().getProperty("principal")!=null)
        && (getProperties().getProperty("keytab")!=null)) {
      try {
        UserGroupInformation.loginUserFromKeytab(getProperties().getProperty("principal"),
              getProperties().getProperty("keytab"));
      } catch (IOException e) {
        System.err.println("Keytab file is not readable or not found");
        throw new DBException(e);
      }
    }

    try {
      threadCount.getAndIncrement();
      synchronized (CONNECTION_LOCK) {
        if (connection == null) {
          // Initialize if not set up already.
          connection = ConnectionFactory.createConnection(config);
        }
      }
    } catch (java.io.IOException e) {
      throw new DBException(e);
    }

    if ((getProperties().getProperty("debug") != null)
        && (getProperties().getProperty("debug").compareTo("true") == 0)) {
      debug = true;
    }

    if ("false"
        .equals(getProperties().getProperty("hbase.usepagefilter", "true"))) {
      usePageFilter = false;
    }

    columnFamily = getProperties().getProperty("columnfamily");
    if (columnFamily == null) {
      System.err.println("Error, must specify a columnfamily for HBase table");
      throw new DBException("No columnfamily specified");
    }
    columnFamilyBytes = Bytes.toBytes(columnFamily);

    // Terminate right now if table does not exist, since the client
    // will not propagate this error upstream once the workload
    // starts.
    String table = com.yahoo.ycsb.workloads.CoreWorkload.table;
    try {
      final TableName tName = TableName.valueOf(table);
      synchronized (CONNECTION_LOCK) {
        connection.getTable(tName).getTableDescriptor();
      }
    } catch (IOException e) {
      throw new DBException(e);
    }
  }

  /**
   * Cleanup any state for this DB. Called once per DB instance; there is one DB
   * instance per client thread.
   */
  @Override
  public void cleanup() throws DBException {
    // Get the measurements instance as this is the only client that should
    // count clean up time like an update if client-side buffering is
    // enabled.
    Measurements measurements = Measurements.getMeasurements();
    try {
      long st = System.nanoTime();
      if (bufferedMutator != null) {
        bufferedMutator.close();
      }
      if (currentTable != null) {
        currentTable.close();
      }
      long en = System.nanoTime();
      final String type = clientSideBuffering ? "UPDATE" : "CLEANUP";
      measurements.measure(type, (int) ((en - st) / 1000));
      threadCount.decrementAndGet();
      if (threadCount.get() <= 0) {
        // Means we are done so ok to shut down the Connection.
        synchronized (CONNECTION_LOCK) {
          if (connection != null) {
            connection.close();
            connection = null;
          }
        }
      }
    } catch (IOException e) {
      throw new DBException(e);
    }
  }

  public void getHTable(String table) throws IOException {
    final TableName tName = TableName.valueOf(table);
    synchronized (CONNECTION_LOCK) {
      this.currentTable = connection.getTable(tName);
      if (clientSideBuffering) {
        final BufferedMutatorParams p = new BufferedMutatorParams(tName);
        p.writeBufferSize(writeBufferSize);
        this.bufferedMutator = connection.getBufferedMutator(p);
      }
    }
  }

  /**
   * 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
   */
  public Status read(String table, String key, Set<String> fields,
      HashMap<String, ByteIterator> result) {
    // if this is a "new" table, init HTable object. Else, use existing one
    if (!tableName.equals(table)) {
      currentTable = null;
      try {
        getHTable(table);
        tableName = table;
      } catch (IOException e) {
        System.err.println("Error accessing HBase table: " + e);
        return Status.ERROR;
      }
    }

    Result r = null;
    try {
      if (debug) {
        System.out
            .println("Doing read from HBase columnfamily " + columnFamily);
        System.out.println("Doing read for key: " + key);
      }
      Get g = new Get(Bytes.toBytes(key));
      if (fields == null) {
        g.addFamily(columnFamilyBytes);
      } else {
        for (String field : fields) {
          g.addColumn(columnFamilyBytes, Bytes.toBytes(field));
        }
      }
      r = currentTable.get(g);
    } catch (IOException e) {
      if (debug) {
        System.err.println("Error doing get: " + e);
      }
      return Status.ERROR;
    } catch (ConcurrentModificationException e) {
      // do nothing for now...need to understand HBase concurrency model better
      return Status.ERROR;
    }

    if (r.isEmpty()) {
      return Status.NOT_FOUND;
    }

    while (r.advance()) {
      final Cell c = r.current();
      result.put(Bytes.toString(CellUtil.cloneQualifier(c)),
          new ByteArrayByteIterator(CellUtil.cloneValue(c)));
      if (debug) {
        System.out.println(
            "Result for field: " + Bytes.toString(CellUtil.cloneQualifier(c))
                + " is: " + Bytes.toString(CellUtil.cloneValue(c)));
      }
    }
    return Status.OK;
  }

  /**
   * 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
   */
  @Override
  public Status scan(String table, String startkey, int recordcount,
      Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
    // if this is a "new" table, init HTable object. Else, use existing one
    if (!tableName.equals(table)) {
      currentTable = null;
      try {
        getHTable(table);
        tableName = table;
      } catch (IOException e) {
        System.err.println("Error accessing HBase table: " + e);
        return Status.ERROR;
      }
    }

    Scan s = new Scan(Bytes.toBytes(startkey));
    // HBase has no record limit. Here, assume recordcount is small enough to
    // bring back in one call.
    // We get back recordcount records
    s.setCaching(recordcount);
    if (this.usePageFilter) {
      s.setFilter(new PageFilter(recordcount));
    }

    // add specified fields or else all fields
    if (fields == null) {
      s.addFamily(columnFamilyBytes);
    } else {
      for (String field : fields) {
        s.addColumn(columnFamilyBytes, Bytes.toBytes(field));
      }
    }

    // get results
    ResultScanner scanner = null;
    try {
      scanner = currentTable.getScanner(s);
      int numResults = 0;
      for (Result rr = scanner.next(); rr != null; rr = scanner.next()) {
        // get row key
        String key = Bytes.toString(rr.getRow());

        if (debug) {
          System.out.println("Got scan result for key: " + key);
        }

        HashMap<String, ByteIterator> rowResult =
            new HashMap<String, ByteIterator>();

        while (rr.advance()) {
          final Cell cell = rr.current();
          rowResult.put(Bytes.toString(CellUtil.cloneQualifier(cell)),
              new ByteArrayByteIterator(CellUtil.cloneValue(cell)));
        }

        // add rowResult to result vector
        result.add(rowResult);
        numResults++;

        // PageFilter does not guarantee that the number of results is <=
        // pageSize, so this
        // break is required.
        if (numResults >= recordcount) {// if hit recordcount, bail out
          break;
        }
      } // done with row
    } catch (IOException e) {
      if (debug) {
        System.out.println("Error in getting/parsing scan result: " + e);
      }
      return Status.ERROR;
    } finally {
      if (scanner != null) {
        scanner.close();
      }
    }

    return Status.OK;
  }

  /**
   * 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
   */
  @Override
  public Status update(String table, String key,
      HashMap<String, ByteIterator> values) {
    // if this is a "new" table, init HTable object. Else, use existing one
    if (!tableName.equals(table)) {
      currentTable = null;
      try {
        getHTable(table);
        tableName = table;
      } catch (IOException e) {
        System.err.println("Error accessing HBase table: " + e);
        return Status.ERROR;
      }
    }

    if (debug) {
      System.out.println("Setting up put for key: " + key);
    }
    Put p = new Put(Bytes.toBytes(key));
    p.setDurability(durability);
    for (Map.Entry<String, ByteIterator> entry : values.entrySet()) {
      byte[] value = entry.getValue().toArray();
      if (debug) {
        System.out.println("Adding field/value " + entry.getKey() + "/"
            + Bytes.toStringBinary(value) + " to put request");
      }
      p.addColumn(columnFamilyBytes, Bytes.toBytes(entry.getKey()), value);
    }

    try {
      if (clientSideBuffering) {
        Preconditions.checkNotNull(bufferedMutator);
        bufferedMutator.mutate(p);
      } else {
        currentTable.put(p);
      }
    } catch (IOException e) {
      if (debug) {
        System.err.println("Error doing put: " + e);
      }
      return Status.ERROR;
    } catch (ConcurrentModificationException e) {
      // do nothing for now...hope this is rare
      return Status.ERROR;
    }

    return Status.OK;
  }

  /**
   * 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
   */
  @Override
  public Status insert(String table, String key,
      HashMap<String, ByteIterator> values) {
    return update(table, key, values);
  }

  /**
   * 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
   */
  @Override
  public Status delete(String table, String key) {
    // if this is a "new" table, init HTable object. Else, use existing one
    if (!tableName.equals(table)) {
      currentTable = null;
      try {
        getHTable(table);
        tableName = table;
      } catch (IOException e) {
        System.err.println("Error accessing HBase table: " + e);
        return Status.ERROR;
      }
    }

    if (debug) {
      System.out.println("Doing delete for key: " + key);
    }

    final Delete d = new Delete(Bytes.toBytes(key));
    d.setDurability(durability);
    try {
      if (clientSideBuffering) {
        Preconditions.checkNotNull(bufferedMutator);
        bufferedMutator.mutate(d);
      } else {
        currentTable.delete(d);
      }
    } catch (IOException e) {
      if (debug) {
        System.err.println("Error doing delete: " + e);
      }
      return Status.ERROR;
    }

    return Status.OK;
  }

  @VisibleForTesting
  void setConfiguration(final Configuration newConfig) {
    this.config = newConfig;
  }
}

/*
 * For customized vim control set autoindent set si set shiftwidth=4
 */
