//  Copyright (c) 2007-2008 Facebook
//
//  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 file LICENSE or visit the Scribe site at:
// http://developers.facebook.com/scribe/
//
// @author Bobby Johnson
// @author James Wang
// @author Jason Sobel
// @author Alex Moskalyuk
// @author Avinash Lakshman
// @author Anthony Giardullo
// @author Jan Oravec
// @author John Song

#include "Common.h"
#include "BucketStore.h"
#include "ScribeServer.h"

using namespace std;

static const char kDefaultBucketStoreDelimiter = ':';

namespace scribe {

BucketStore::BucketStore(StoreQueue* storeq,
                        const string& category,
                        bool multiCategory)
  : Store(storeq, category, "bucket", multiCategory),
    bucketType_(CONTEXT_LOG),
    delimiter_(kDefaultBucketStoreDelimiter),
    removeKey_(false),
    opened_(false),
    bucketRange_(0),
    numBuckets_(1) {
}

BucketStore::~BucketStore() {

}

// Given a single bucket definition, create multiple buckets
void BucketStore::createBucketsFromBucket(StoreConfPtr configuration,
    StoreConfPtr bucketConf) {
  string errorMesg, bucketSubdir, type, path, failureBucket;
  bool needsBucketSubdir = false;
  unsigned long bucketOffset = 0;
  StoreConfPtr tmp;

  // check for extra bucket definitions
  if (configuration->getStore("bucket0", &tmp) ||
      configuration->getStore("bucket1", &tmp)) {
    errorMesg = "bucket store has too many buckets defined";
    goto handle_error;
  }

  bucketConf->getString("type", &type);
  if (type != "file" && type != "thriftfile") {
    errorMesg = "store contained in a bucket store must have a type of ";
    errorMesg += "either file or thriftfile if not defined explicitely";
    goto handle_error;
  }

  needsBucketSubdir = true;
  if (!configuration->getString("bucket_subdir", &bucketSubdir)) {
    errorMesg = "bucketizer containing file stores must have a bucket_subdir";
    goto handle_error;
  }
  if (!bucketConf->getString("file_path", &path)) {
    errorMesg = "file store contained by bucketizer must have a file_path";
    goto handle_error;
  }

  // set starting bucket number if specified
  configuration->getUnsigned("bucket_offset", &bucketOffset);

  // check if failure bucket was given a different name
  configuration->getString("failure_bucket", &failureBucket);

  // We actually create numBuckets_ + 1 stores. Messages are normally
  // hashed into buckets 1 through numBuckets_, and messages that can't
  // be hashed are put in bucket 0.

  for (unsigned int i = 0; i <= numBuckets_; ++i) {

    StorePtr newStore =
      createStore(storeQueue_, type, categoryHandled_, false, multiCategory_);

    if (!newStore) {
      errorMesg = "can't create store of type: ";
      errorMesg += type;
      goto handle_error;
    }

    // For file/thrift file buckets, create unique filepath for each bucket
    if (needsBucketSubdir) {
      if (i == 0 && !failureBucket.empty()) {
        bucketConf->setString("file_path", path + '/' + failureBucket);
      } else {
        // the bucket number is appended to the file path
        unsigned int bucketId = i + bucketOffset;

        ostringstream oss;
        oss << path << '/' << bucketSubdir << setw(3) << setfill('0')
            << bucketId;
        bucketConf->setString("file_path", oss.str());
      }
    }

    buckets_.push_back(newStore);
    newStore->configure(bucketConf, storeConf_);
  }

  return;

handle_error:
  setStatus(errorMesg);
  LOG_OPER("[%s] Bad config - %s", categoryHandled_.c_str(),
           errorMesg.c_str());
  numBuckets_ = 0;
  buckets_.clear();
}

// Checks for a bucket definition for every bucket from 0 to numBuckets_
// and configures each bucket
void BucketStore::createBuckets(StoreConfPtr configuration) {
  string errorMesg, tmpString;
  StoreConfPtr tmp;

  if (configuration->getString("bucket_subdir", &tmpString)) {
    errorMesg = "cannot have bucket_subdir when defining multiple buckets";
    goto handle_error;
  }

  if (configuration->getString("bucket_offset", &tmpString)) {
    errorMesg = "cannot have bucket_offset when defining multiple buckets";
    goto handle_error;
  }

  if (configuration->getString("failure_bucket", &tmpString)) {
    errorMesg = "cannot have failure_bucket when defining multiple buckets";
    goto handle_error;
  }

  // Configure stores named 'bucket0, bucket1, bucket2, ... bucket{numBuckets_}
  for (unsigned long i = 0; i <= numBuckets_; i++) {
    StoreConfPtr   bucketConf;
    string       type, bucketName;
    stringstream ss;

    ss << "bucket" << i;
    bucketName = ss.str();

    if (!configuration->getStore(bucketName, &bucketConf)) {
      errorMesg = "could not find bucket definition for " + bucketName;
      goto handle_error;
    }

    if (!bucketConf->getString("type", &type)) {
      errorMesg = "store contained in a bucket store must have a type";
      goto handle_error;
    }

    StorePtr bucket =
      createStore(storeQueue_, type, categoryHandled_, false, multiCategory_);

    buckets_.push_back(bucket);
    //add bucket id configuration
    bucketConf->setUnsigned("bucket_id", i);
    bucketConf->setUnsigned("network::bucket_id", i);
    bucketConf->setUnsigned("file::bucket_id", i);
    bucketConf->setUnsigned("thriftfile::bucket_id", i);
    bucketConf->setUnsigned("buffer::bucket_id", i);
    bucket->configure(bucketConf, storeConf_);
  }

  // Check if an extra bucket is defined
  if (configuration->getStore("bucket" + (numBuckets_ + 1), &tmp)) {
    errorMesg = "bucket store has too many buckets defined";
    goto handle_error;
  }

  return;

handle_error:
  setStatus(errorMesg);
  LOG_OPER("[%s] Bad config - %s", categoryHandled_.c_str(),
           errorMesg.c_str());
  numBuckets_ = 0;
  buckets_.clear();
}

/**
   * Buckets in a bucket store can be defined explicitly or implicitly:
   *
   * #Explicitly
   * <store>
   *   type=bucket
   *   num_buckets=2
   *   bucket_type=key_hash
   *
   *   <bucket0>
   *     ...
   *   </bucket0>
   *
   *   <bucket1>
   *     ...
   *   </bucket1>
   *
   *   <bucket2>
   *     ...
   *   </bucket2>
   * </store>
   *
   * #Implicitly
   * <store>
   *   type=bucket
   *   num_buckets=2
   *   bucket_type=key_hash
   *
   *   <bucket>
   *     ...
   *   </bucket>
   * </store>
   */
void BucketStore::configure(StoreConfPtr configuration, StoreConfPtr parent) {
  Store::configure(configuration, parent);

  string errorMesg, bucketizerStr, removeKeyStr;
  unsigned long delimLong = 0;
  StoreConfPtr bucketConf;
  //set this to true for bucket types that have a delimiter
  bool needDelimiter = false;

  configuration->getString("bucket_type", &bucketizerStr);

  // Figure out th bucket type from the bucketizer string
  if (0 == bucketizerStr.compare("context_log")) {
    bucketType_ = CONTEXT_LOG;
  } else if (0 == bucketizerStr.compare("random")) {
      bucketType_ = RANDOM;
  } else if (0 == bucketizerStr.compare("key_hash")) {
    bucketType_ = KEY_HASH;
    needDelimiter = true;
  } else if (0 == bucketizerStr.compare("key_modulo")) {
    bucketType_ = KEY_MODULO;
    needDelimiter = true;
  } else if (0 == bucketizerStr.compare("key_range")) {
    bucketType_ = KEY_RANGE;
    needDelimiter = true;
    configuration->getUnsigned("bucket_range", &bucketRange_);

    if (bucketRange_ == 0) {
      LOG_OPER("[%s] config warning - bucket_range is 0",
               categoryHandled_.c_str());
    }
  }

  // This is either a key_hash or key_modulo, not context log, figure out
  // the delimiter and store it
  if (needDelimiter) {
    configuration->getUnsigned("delimiter", &delimLong);
    if (delimLong > 255) {
      LOG_OPER("[%s] config warning - delimiter is too large to fit in a char, "
               "using default", categoryHandled_.c_str());
      delimiter_ = kDefaultBucketStoreDelimiter;
    } else if (delimLong == 0) {
      LOG_OPER("[%s] config warning - delimiter is zero, using default",
               categoryHandled_.c_str());
      delimiter_ = kDefaultBucketStoreDelimiter;
    } else {
      delimiter_ = (char)delimLong;
    }
  }

  // Optionally remove the key and delimiter of each message before bucketizing
  configuration->getString("remove_key", &removeKeyStr);
  if (removeKeyStr == "yes") {
    removeKey_ = true;

    if (bucketType_ == CONTEXT_LOG) {
      errorMesg = "Bad config - bucketizer store of type context_log do not "
                  "support remove_key";
      goto handle_error;
    }
  }

  if (!configuration->getUnsigned("num_buckets", &numBuckets_)) {
    errorMesg = "Bad config - bucket store must have num_buckets";
    goto handle_error;
  }

  // Buckets can be defined explicitely or by specifying a single "bucket"
  if (configuration->getStore("bucket", &bucketConf)) {
    createBucketsFromBucket(configuration, bucketConf);
  } else {
    createBuckets(configuration);
  }

  return;

handle_error:
  setStatus(errorMesg);
  LOG_OPER("[%s] %s", categoryHandled_.c_str(), errorMesg.c_str());
  numBuckets_ = 0;
  buckets_.clear();
}

bool BucketStore::open() {
  // we have one extra bucket for messages we can't hash
  if (numBuckets_ <= 0 || buckets_.size() != numBuckets_ + 1) {
    LOG_OPER("[%s] Can't open bucket store with <%d> of <%lu> buckets",
             categoryHandled_.c_str(), (int)buckets_.size(), numBuckets_);
    return false;
  }

  for (vector<StorePtr>::iterator iter = buckets_.begin();
       iter != buckets_.end();
       ++iter) {

    if (!(*iter)->open()) {
      close();
      opened_ = false;
      return false;
    }
  }
  opened_ = true;
  return true;
}

bool BucketStore::isOpen() {
  return opened_;
}

void BucketStore::close() {
  // don't check opened, because we can call this when some, but
  // not all, contained stores are opened. Calling close on a contained
  // store that's already closed shouldn't hurt anything.
  for (vector<StorePtr>::iterator iter = buckets_.begin();
       iter != buckets_.end();
       ++iter) {
    (*iter)->close();
  }
  opened_ = false;
}

void BucketStore::flush() {
  for (vector<StorePtr>::iterator iter = buckets_.begin();
       iter != buckets_.end();
       ++iter) {
    (*iter)->flush();
  }
}

string BucketStore::getStatus() {

  string retVal = Store::getStatus();

  vector<StorePtr>::iterator iter = buckets_.begin();
  while (retVal.empty() && iter != buckets_.end()) {
    retVal = (*iter)->getStatus();
    ++iter;
  }
  return retVal;
}

// Call periodicCheck on all containing stores
void BucketStore::periodicCheck() {
  // Call periodic check on all bucket stores in a random order
  uint32_t sz = buckets_.size();
  vector<uint32_t> storeIndex(sz);
  for (uint32_t i = 0; i < sz; ++i) {
    storeIndex[i] = i;
  }
  random_shuffle(storeIndex.begin(), storeIndex.end());

  for (uint32_t i = 0; i < sz; ++i) {
    uint32_t idx = storeIndex[i];
    buckets_[idx]->periodicCheck();
  }
}

StorePtr BucketStore::copy(const string &category) {
  StorePtr copied(new BucketStore(storeQueue_, category, multiCategory_));

  BucketStore* store = static_cast<BucketStore*>(copied.get());
  store->numBuckets_ = numBuckets_;
  store->bucketType_ = bucketType_;
  store->delimiter_ = delimiter_;

  for (vector<StorePtr>::iterator iter = buckets_.begin();
       iter != buckets_.end();
       ++iter) {
    store->buckets_.push_back((*iter)->copy(category));
  }

  return copied;
}

/*
 * Bucketize <messages> and try to send to each contained bucket store
 * At the end of the function <messages> will contain all the messages that
 * could not be processed
 * Returns true if all messages were successfully sent, false otherwise.
 */
bool BucketStore::handleMessages(LogEntryVectorPtr messages) {
  bool success = true;

  LogEntryVectorPtr failedMessages(new LogEntryVector);
  vector<LogEntryVectorPtr> bucketedMessages;
  // bucket numbers are 1-based while bucket #0 is reserved for
  // error handling
  bucketedMessages.resize(numBuckets_ + 1);

  if (numBuckets_ == 0) {
    LOG_OPER("[%s] Failed to write - no buckets configured",
             categoryHandled_.c_str());
    setStatus("Failed write to bucket store");
    return false;
  }

  // batch messages by bucket
  for (LogEntryVector::iterator iter = messages->begin();
       iter != messages->end();
       ++iter) {
    unsigned bucket = bucketize((*iter)->message);

    if (!bucketedMessages[bucket]) {
      bucketedMessages[bucket].reset(new LogEntryVector);
    }

    bucketedMessages[bucket]->push_back(*iter);
  }

  // handle all batches of messages
  for (unsigned long i = 0; i <= numBuckets_; i++) {
    LogEntryVectorPtr batch = bucketedMessages[i];

    if (batch) {

      if (removeKey_) {
        // Create new set of messages with keys removed
        LogEntryVectorPtr keyRemoved(new LogEntryVector);

        for (LogEntryVector::iterator iter = batch->begin();
             iter != batch->end();
             ++iter) {
          const LogEntryPtr& original = *iter;

          LogEntryPtr entry(new LogEntry);
          *entry = *original;
          // strip the key from message
          entry->message = getMessageWithoutKey(entry->message);
          keyRemoved->push_back(entry);
        }
        batch = keyRemoved;
      }

      if (!buckets_[i]->handleMessages(batch)) {
        // keep track of messages that were not handled
        failedMessages->insert(failedMessages->end(),
                               bucketedMessages[i]->begin(),
                               bucketedMessages[i]->end());
        success = false;
      }
    }
  }

  if (!success) {
    // return failed logentrys in messages
    messages->swap(*failedMessages);
  }

  return success;
}

// Return the bucket number a message must be put into;
// Return 0 on any errors
unsigned long BucketStore::bucketize(const string& message) {

  string::size_type length = message.length();

  if (bucketType_ == CONTEXT_LOG) {
    // the key is in ascii after the third delimiter
    char delim = 1;
    string::size_type pos = 0;
    for (int i = 0; i < 3; ++i) {
      pos = message.find(delim, pos);
      if (pos == string::npos || length <= pos + 1) {
        return 0;
      }
      ++pos;
    }
    if (message[pos] == delim) {
      return 0;
    }

    uint32_t id = strtoul(message.substr(pos).c_str(), NULL, 10);
    if (id == 0) {
      return 0;
    }

    if (numBuckets_ == 0) {
      return 0;
    } else {
      return (scribe::integerhash::hash32(id) % numBuckets_) + 1;
    }
  } else if (bucketType_ == RANDOM) {
    // return any random bucket
    return (rand() % numBuckets_) + 1;
  } else {
    // just hash everything before the first user-defined delimiter
    string::size_type pos = message.find(delimiter_);
    if (pos == string::npos) {
      // if no delimiter found, write to bucket 0
      return 0;
    }

    string key = message.substr(0, pos).c_str();
    if (key.empty()) {
      // if no key found, write to bucket 0
      return 0;
    }

    if (numBuckets_ == 0) {
      return 0;
    } else {
      switch (bucketType_) {
        case KEY_MODULO:
          // No hashing, just simple modulo
          return (atol(key.c_str()) % numBuckets_) + 1;
          break;
        case KEY_RANGE:
          if (bucketRange_ == 0) {
            return 0;
          } else {
            // Calculate what bucket this key would fall into if we used
            // bucket_range to compute the modulo
           double keyMod = atol(key.c_str()) % bucketRange_;
           return (unsigned long) ((keyMod / bucketRange_) * numBuckets_) + 1;
          }
          break;
        case KEY_HASH:
        default:
          // Hashing by default.
          return (scribe::strhash::hash32(key.c_str()) % numBuckets_) + 1;
          break;
      }
    }
  }

  return 0;
}

string BucketStore::getMessageWithoutKey(const string& message) {
  string::size_type pos = message.find(delimiter_);

  if (pos == string::npos) {
    return message;
  }

  return message.substr(pos+1);
}

} //! namespace scribe