/**
 * Copyright (c) 2016 YCSB contributors. All rights reserved.
 *
 * 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.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.Status;
import com.yahoo.ycsb.StringByteIterator;

import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.impl.CloudSolrClient;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.client.solrj.response.UpdateResponse;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.SolrInputDocument;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;

/**
 * Solr client for YCSB framework.
 *
 * <p>
 * Default properties to set:
 * </p>
 * <ul>
 * See README.md
 * </ul>
 *
 */
public class SolrClient extends DB {

  public static final String DEFAULT_CLOUD_MODE = "false";
  public static final String DEFAULT_BATCH_MODE = "false";
  public static final String DEFAULT_ZOOKEEPER_HOSTS = "localhost:2181";
  public static final String DEFAULT_SOLR_BASE_URL = "http://localhost:8983/solr";
  public static final String DEFAULT_COMMIT_WITHIN_TIME = "1000";

  private org.apache.solr.client.solrj.SolrClient client;
  private Integer commitTime;
  private Boolean batchMode;

  /**
   * 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 {
    Properties props = getProperties();
    commitTime = Integer
        .parseInt(props.getProperty("solr.commit.within.time", DEFAULT_COMMIT_WITHIN_TIME));
    batchMode = Boolean.parseBoolean(props.getProperty("solr.batch.mode", DEFAULT_BATCH_MODE));

    // Check if Solr cluster is running in SolrCloud or Stand-alone mode
    Boolean cloudMode = Boolean.parseBoolean(props.getProperty("solr.cloud", DEFAULT_CLOUD_MODE));
    System.err.println("Solr Cloud Mode = " + cloudMode);
    if (cloudMode) {
      System.err.println("Solr Zookeeper Remote Hosts = "
          + props.getProperty("solr.zookeeper.hosts", DEFAULT_ZOOKEEPER_HOSTS));
      client = new CloudSolrClient(
          props.getProperty("solr.zookeeper.hosts", DEFAULT_ZOOKEEPER_HOSTS));
    } else {
      client = new HttpSolrClient(props.getProperty("solr.base.url", DEFAULT_SOLR_BASE_URL));
    }
  }

  @Override
  public void cleanup() throws DBException {
    try {
      client.close();
    } catch (IOException e) {
      throw new DBException(e);
    }
  }

  /**
   * 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.
   */
  @Override
  public Status insert(String table, String key, HashMap<String, ByteIterator> values) {
    try {
      SolrInputDocument doc = new SolrInputDocument();

      doc.addField("id", key);
      for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
        doc.addField(entry.getKey(), entry.getValue());
      }
      UpdateResponse response;
      if (batchMode) {
        response = client.add(table, doc, commitTime);
      } else {
        response = client.add(table, doc);
        client.commit(table);
      }
      return checkStatus(response.getStatus());

    } catch (IOException | SolrServerException e) {
      e.printStackTrace();
    }
    return Status.ERROR;
  }

  /**
   * 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.
   */
  @Override
  public Status delete(String table, String key) {
    try {
      UpdateResponse response;
      if (batchMode) {
        response = client.deleteById(table, key, commitTime);
      } else {
        response = client.deleteById(table, key);
        client.commit(table);
      }
      return checkStatus(response.getStatus());
    } catch (IOException | SolrServerException e) {
      e.printStackTrace();
    }
    return Status.ERROR;
  }

  /**
   * 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".
   */
  @Override
  public Status read(String table, String key, Set<String> fields,
      HashMap<String, ByteIterator> result) {
    try {
      Boolean returnFields = false;
      String[] fieldList = null;
      if (fields != null) {
        returnFields = true;
        fieldList = fields.toArray(new String[fields.size()]);
      }
      SolrQuery query = new SolrQuery();
      query.setQuery("id:" + key);
      if (returnFields) {
        query.setFields(fieldList);
      }
      final QueryResponse response = client.query(table, query);
      SolrDocumentList results = response.getResults();
      if ((results != null) && (results.getNumFound() > 0)) {
        for (String field : results.get(0).getFieldNames()) {
          result.put(field,
              new StringByteIterator(String.valueOf(results.get(0).getFirstValue(field))));
        }
      }
      return checkStatus(response.getStatus());
    } catch (IOException | SolrServerException e) {
      e.printStackTrace();
    }
    return Status.ERROR;
  }

  /**
   * 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.
   */
  @Override
  public Status update(String table, String key, HashMap<String, ByteIterator> values) {
    try {
      SolrInputDocument updatedDoc = new SolrInputDocument();
      updatedDoc.addField("id", key);

      for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
        updatedDoc.addField(entry.getKey(), Collections.singletonMap("set", entry.getValue()));
      }

      UpdateResponse writeResponse;
      if (batchMode) {
        writeResponse = client.add(table, updatedDoc, commitTime);
      } else {
        writeResponse = client.add(table, updatedDoc);
        client.commit(table);
      }
      return checkStatus(writeResponse.getStatus());
    } catch (IOException | SolrServerException e) {
      e.printStackTrace();
    }
    return Status.ERROR;
  }

  /**
   * 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.
   */
  @Override
  public Status scan(String table, String startkey, int recordcount, Set<String> fields,
      Vector<HashMap<String, ByteIterator>> result) {
    try {
      Boolean returnFields = false;
      String[] fieldList = null;
      if (fields != null) {
        returnFields = true;
        fieldList = fields.toArray(new String[fields.size()]);
      }
      SolrQuery query = new SolrQuery();
      query.setQuery("*:*");
      query.setParam("fq", "id:[ " + startkey + " TO * ]");
      if (returnFields) {
        query.setFields(fieldList);
      }
      query.setRows(recordcount);
      final QueryResponse response = client.query(table, query);
      SolrDocumentList results = response.getResults();

      HashMap<String, ByteIterator> entry;

      for (SolrDocument hit : results) {
        entry = new HashMap<String, ByteIterator>((int) results.getNumFound());
        for (String field : hit.getFieldNames()) {
          entry.put(field, new StringByteIterator(String.valueOf(hit.getFirstValue(field))));
        }
        result.add(entry);
      }
      return checkStatus(response.getStatus());
    } catch (IOException | SolrServerException e) {
      e.printStackTrace();
    }
    return Status.ERROR;
  }

  private Status checkStatus(int status) {
    Status responseStatus;
    switch (status) {
    case 0:
      responseStatus = Status.OK;
      break;
    case 400:
      responseStatus = Status.BAD_REQUEST;
      break;
    case 403:
      responseStatus = Status.FORBIDDEN;
      break;
    case 404:
      responseStatus = Status.NOT_FOUND;
      break;
    case 500:
      responseStatus = Status.ERROR;
      break;
    case 503:
      responseStatus = Status.SERVICE_UNAVAILABLE;
      break;
    default:
      responseStatus = Status.UNEXPECTED_STATE;
      break;
    }
    return responseStatus;
  }

}
