/*
 * librdkafka - Apache Kafka C library
 *
 * Copyright (c) 2012-2018, Magnus Edenhill
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <iostream>
#include <fstream>
#include <iterator>
#include <string>
#include "testcpp.h"

#if WITH_RAPIDJSON
#include <rapidjson/document.h>
#include <rapidjson/schema.h>
#include <rapidjson/filereadstream.h>
#include <rapidjson/stringbuffer.h>
#include <rapidjson/error/en.h>
#include <rapidjson/prettywriter.h>
#endif

static const char *stats_schema_path = "../src/statistics_schema.json";

#if WITH_RAPIDJSON
/**
 * @brief Statistics schema validator
 */
class TestSchemaValidator {
 public:
  TestSchemaValidator () {

  }
  TestSchemaValidator (const std::string schema_path) {
    /* Read schema from file */
    schema_path_ = schema_path;

    std::ifstream f(schema_path.c_str());
    if (!f.is_open())
            Test::Fail(tostr() << "Failed to open schema " << schema_path <<
                       ": " << strerror(errno));
    std::string schema_str((std::istreambuf_iterator<char>(f)),
                           (std::istreambuf_iterator<char>()));

    /* Parse schema */
    sd_ = new rapidjson::Document();
    if (sd_->Parse(schema_str.c_str()).HasParseError())
      Test::Fail(tostr() << "Failed to parse statistics schema: " <<
                 rapidjson::GetParseError_En(sd_->GetParseError()) <<
                 " at " << sd_->GetErrorOffset());

    schema_ = new rapidjson::SchemaDocument(*sd_);
    validator_ = new rapidjson::SchemaValidator(*schema_);
  }

  ~TestSchemaValidator () {
    if (sd_)
      delete sd_;
    if (schema_)
      delete schema_;
    if (validator_)
      delete validator_;
  }

  void validate (const std::string &json_doc) {
    /* Parse JSON to validate */
    rapidjson::Document d;
    if (d.Parse(json_doc.c_str()).HasParseError())
      Test::Fail(tostr() << "Failed to parse stats JSON: " <<
                 rapidjson::GetParseError_En(d.GetParseError()) <<
                 " at " << d.GetErrorOffset());

    /* Validate using schema */
    if (!d.Accept(*validator_)) {

      rapidjson::StringBuffer sb;

      validator_->GetInvalidSchemaPointer().StringifyUriFragment(sb);
      Test::Say(tostr() << "Schema: " << sb.GetString() << "\n");
      Test::Say(tostr() << "Invalid keyword: " << validator_->GetInvalidSchemaKeyword() << "\n");
      sb.Clear();

      validator_->GetInvalidDocumentPointer().StringifyUriFragment(sb);
      Test::Say(tostr() << "Invalid document: " << sb.GetString() << "\n");
      sb.Clear();

      Test::Fail(tostr() << "JSON validation using schema " << schema_path_ << " failed");
    }

    Test::Say(3, "JSON document validated using schema " + schema_path_ + "\n");
  }

 private:
  std::string schema_path_;
  rapidjson::Document *sd_;
  rapidjson::SchemaDocument *schema_;
  rapidjson::SchemaValidator *validator_;
};


#else

/* Dummy validator doing nothing when RapidJSON is unavailable */
class TestSchemaValidator {
 public:
  TestSchemaValidator () {

  }
  TestSchemaValidator (const std::string schema_path) {
  }

  ~TestSchemaValidator () {
  }

  void validate (const std::string &json_doc) {
  }
};

#endif

class myEventCb : public RdKafka::EventCb {
 public:
  myEventCb(const std::string schema_path):
      validator_(TestSchemaValidator(schema_path)) {
    stats_cnt = 0;
  }

  int stats_cnt;
  std::string last;  /**< Last stats document */

  void event_cb (RdKafka::Event &event) {
    switch (event.type())
    {
      case RdKafka::Event::EVENT_STATS:
        if (!(stats_cnt % 10))
          Test::Say(tostr() << "Stats (#" << stats_cnt << "): " <<
                    event.str() << "\n");
        if (event.str().length() > 20)
          stats_cnt += 1;
        validator_.validate(event.str());
        last = event.str();
        break;
      default:
        break;
    }
  }

 private:
  TestSchemaValidator validator_;
};


/**
 * @brief Verify that stats are emitted according to statistics.interval.ms
 */
void test_stats_timing () {
  RdKafka::Conf *conf = RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL);
  myEventCb my_event = myEventCb(stats_schema_path);
  std::string errstr;

  if (conf->set("statistics.interval.ms", "100", errstr) != RdKafka::Conf::CONF_OK)
          Test::Fail(errstr);

  if (conf->set("event_cb", &my_event, errstr) != RdKafka::Conf::CONF_OK)
          Test::Fail(errstr);

  RdKafka::Producer *p = RdKafka::Producer::create(conf, errstr);
  if (!p)
          Test::Fail("Failed to create Producer: " + errstr);
  delete conf;

  int64_t t_start = test_clock();

  while (my_event.stats_cnt < 12)
    p->poll(1000);

  int elapsed = (int)((test_clock() - t_start) / 1000);
  const int expected_time = 1200;

  Test::Say(tostr() << my_event.stats_cnt << " (expected 12) stats callbacks received in " <<
            elapsed << "ms (expected " << expected_time << "ms +-25%)\n");

  if (elapsed < expected_time * 0.75 ||
      elapsed > expected_time * 1.25) {
    /* We can't rely on CIs giving our test job enough CPU to finish
     * in time, so don't error out even if the time is outside the window */
    if (test_on_ci)
      Test::Say(tostr() << "WARNING: Elapsed time " << elapsed << "ms outside +-25% window (" <<
                expected_time << "ms), cnt " << my_event.stats_cnt);
    else
      Test::Fail(tostr() << "Elapsed time " << elapsed << "ms outside +-25% window (" <<
                 expected_time << "ms), cnt " << my_event.stats_cnt);
  }
  delete p;
}



#if WITH_RAPIDJSON

/**
 * @brief Expected partition stats
 */
struct exp_part_stats {
  std::string topic;   /**< Topic */
  int32_t part;        /**< Partition id */
  int     msgcnt;      /**< Expected message count */
  int     msgsize;     /**< Expected per message size.
                        *   This includes both key and value lengths */

  /* Calculated */
  int64_t totsize;     /**< Message size sum */
};

/**
 * @brief Verify end-to-end producer and consumer stats.
 */
static void verify_e2e_stats (const std::string &prod_stats,
                              const std::string &cons_stats,
                              struct exp_part_stats *exp_parts, int partcnt) {
  /**
   * Parse JSON stats
   * These documents are already validated in the Event callback.
   */
  rapidjson::Document p;
  if (p.Parse<rapidjson::kParseValidateEncodingFlag>(prod_stats.c_str()).HasParseError())
    Test::Fail(tostr() << "Failed to parse producer stats JSON: " <<
               rapidjson::GetParseError_En(p.GetParseError()) <<
               " at " << p.GetErrorOffset());

  rapidjson::Document c;
  if (c.Parse<rapidjson::kParseValidateEncodingFlag>(cons_stats.c_str()).HasParseError())
    Test::Fail(tostr() << "Failed to parse consumer stats JSON: " <<
               rapidjson::GetParseError_En(c.GetParseError()) <<
               " at " << c.GetErrorOffset());

  assert(p.HasMember("name"));
  assert(c.HasMember("name"));
  assert(p.HasMember("type"));
  assert(c.HasMember("type"));

  Test::Say(tostr() << "Verifying stats from Producer " << p["name"].GetString() <<
            " and Consumer " << c["name"].GetString() << "\n");

  assert(!strcmp(p["type"].GetString(), "producer"));
  assert(!strcmp(c["type"].GetString(), "consumer"));

  int64_t exp_tot_txmsgs = 0;
  int64_t exp_tot_txmsg_bytes = 0;
  int64_t exp_tot_rxmsgs = 0;
  int64_t exp_tot_rxmsg_bytes = 0;

  for (int part = 0 ; part < partcnt ; part++) {

    /*
     * Find partition stats.
     */

    /* Construct the partition path. */
    char path[256];
    rd_snprintf(path, sizeof(path),
                "/topics/%s/partitions/%d",
                exp_parts[part].topic.c_str(), exp_parts[part].part);
    Test::Say(tostr() << "Looking up partition " << exp_parts[part].part <<
              " with path " << path << "\n");

    /* Even though GetValueByPointer() takes a "char[]" it can only be used
     * with perfectly sized char buffers or string literals since it
     * does not respect NUL terminators.
     * So instead convert the path to a Pointer.*/
    rapidjson::Pointer jpath((const char *)path);

    rapidjson::Value *pp = rapidjson::GetValueByPointer(p, jpath);
    if (!pp)
      Test::Fail(tostr() << "Producer: could not find " << path <<
                 " in " << prod_stats << "\n");

    rapidjson::Value *cp = rapidjson::GetValueByPointer(c, jpath);
    if (!pp)
      Test::Fail(tostr() << "Consumer: could not find " << path <<
                 " in " << cons_stats << "\n");

    assert(pp->HasMember("partition"));
    assert(pp->HasMember("txmsgs"));
    assert(pp->HasMember("txbytes"));

    assert(cp->HasMember("partition"));
    assert(cp->HasMember("rxmsgs"));
    assert(cp->HasMember("rxbytes"));

    Test::Say(tostr() << "partition: " << (*pp)["partition"].GetInt() << "\n");

    int64_t txmsgs = (*pp)["txmsgs"].GetInt();
    int64_t txbytes = (*pp)["txbytes"].GetInt();
    int64_t rxmsgs = (*cp)["rxmsgs"].GetInt();
    int64_t rxbytes = (*cp)["rxbytes"].GetInt();

    exp_tot_txmsgs += txmsgs;
    exp_tot_txmsg_bytes += txbytes;
    exp_tot_rxmsgs += rxmsgs;
    exp_tot_rxmsg_bytes += rxbytes;

    Test::Say(tostr() << "Producer partition: " << (*pp)["partition"].GetInt() << ": " <<
              "txmsgs: " << txmsgs << " vs " << exp_parts[part].msgcnt << ", " <<
              "txbytes: " << txbytes << " vs " << exp_parts[part].totsize << "\n");
    Test::Say(tostr() << "Consumer partition: " << (*cp)["partition"].GetInt() << ": " <<
              "rxmsgs: " << rxmsgs << " vs " << exp_parts[part].msgcnt << ", " <<
              "rxbytes: " << rxbytes << " vs " << exp_parts[part].totsize << "\n");
  }

  /* Check top-level total stats */

  assert(p.HasMember("txmsgs"));
  assert(p.HasMember("txmsg_bytes"));
  assert(p.HasMember("rxmsgs"));
  assert(p.HasMember("rxmsg_bytes"));

  int64_t tot_txmsgs = p["txmsgs"].GetInt();
  int64_t tot_txmsg_bytes = p["txmsg_bytes"].GetInt();
  int64_t tot_rxmsgs = c["rxmsgs"].GetInt();
  int64_t tot_rxmsg_bytes = c["rxmsg_bytes"].GetInt();

  Test::Say(tostr() << "Producer total: " <<
            "txmsgs: " << tot_txmsgs << " vs " << exp_tot_txmsgs << ", " <<
            "txbytes: " << tot_txmsg_bytes << " vs " << exp_tot_txmsg_bytes << "\n");
  Test::Say(tostr() << "Consumer total: " <<
            "rxmsgs: " << tot_rxmsgs << " vs " << exp_tot_rxmsgs << ", " <<
            "rxbytes: " << tot_rxmsg_bytes << " vs " << exp_tot_rxmsg_bytes << "\n");

}

/**
 * @brief Verify stats JSON structure and individual metric fields.
 *
 * To capture as much verifiable data as possible we run a full
 * producer - consumer end to end test and verify that counters
 * and states are emitted accordingly.
 *
 * Requires RapidJSON (for parsing the stats).
 */
static void test_stats () {
  std::string errstr;
  RdKafka::Conf *conf;
  myEventCb producer_event(stats_schema_path);
  myEventCb consumer_event(stats_schema_path);

  std::string topic = Test::mk_topic_name("0053_stats", 1);

  const int partcnt = 2;
  int msgcnt = (test_quick ? 10 : 100) * partcnt;
  const int msgsize = 6*1024;

  /*
   * Common config for producer and consumer
   */
  Test::conf_init(&conf, NULL, 60);
  if (conf->set("statistics.interval.ms", "1000", errstr) != RdKafka::Conf::CONF_OK)
          Test::Fail(errstr);


  /*
   * Create Producer
   */
  if (conf->set("event_cb", &producer_event, errstr) != RdKafka::Conf::CONF_OK)
          Test::Fail(errstr);

  RdKafka::Producer *p = RdKafka::Producer::create(conf, errstr);
  if (!p)
          Test::Fail("Failed to create Producer: " + errstr);


  /*
   * Create Consumer
   */
  conf->set("group.id", topic, errstr);
  conf->set("auto.offset.reset", "earliest", errstr);
  conf->set("enable.partition.eof", "false", errstr);
  if (conf->set("event_cb", &consumer_event, errstr) != RdKafka::Conf::CONF_OK)
          Test::Fail(errstr);

  RdKafka::KafkaConsumer *c = RdKafka::KafkaConsumer::create(conf, errstr);
  if (!c)
    Test::Fail("Failed to create KafkaConsumer: " + errstr);
  delete conf;

  /*
   * Set up consumer assignment (but assign after producing
   * since there will be no topics now) and expected partitions
   * for later verification.
   */
  std::vector<RdKafka::TopicPartition*> toppars;
  struct exp_part_stats exp_parts[partcnt] = {};

  for (int32_t part = 0 ; part < (int32_t)partcnt ; part++) {
    toppars.push_back(RdKafka::TopicPartition::create(topic, part,
                                                      RdKafka::Topic::OFFSET_BEGINNING));
    exp_parts[part].topic = topic;
    exp_parts[part].part = part;
    exp_parts[part].msgcnt = msgcnt / partcnt;
    exp_parts[part].msgsize = msgsize;
    exp_parts[part].totsize = 0;
  }

  /*
   * Produce messages
   */
  uint64_t testid = test_id_generate();

  char key[256];
  char *buf = (char *)malloc(msgsize);

  for (int32_t part = 0 ; part < (int32_t)partcnt ; part++) {
    for (int i = 0 ; i < msgcnt / partcnt ; i++) {
      test_prepare_msg(testid, part, i, buf, msgsize, key, sizeof(key));
      RdKafka::ErrorCode err = p->produce(topic, part,
                                          RdKafka::Producer::RK_MSG_COPY,
                                          buf, msgsize, key, sizeof(key),
                                          -1, NULL);
      if (err)
        Test::Fail("Produce failed: " + RdKafka::err2str(err));
      exp_parts[part].totsize += msgsize + sizeof(key);
      p->poll(0);
    }
  }

  free(buf);

  Test::Say("Waiting for final message delivery\n");
  /* Wait for delivery */
  p->flush(15*1000);

  /*
  * Start consuming partitions
  */
  c->assign(toppars);
  RdKafka::TopicPartition::destroy(toppars);

  /*
   * Consume the messages
   */
  int recvcnt = 0;
  Test::Say(tostr() << "Consuming " << msgcnt << " messages\n");
  while (recvcnt < msgcnt) {
    RdKafka::Message *msg = c->consume(-1);
    if (msg->err())
      Test::Fail("Consume failed: " + msg->errstr());

    int msgid;
    TestMessageVerify(testid, -1, &msgid, msg);
    recvcnt++;
    delete msg;
  }

  /*
   * Producer:
   * Wait for one last stats emit when all messages have been delivered.
   */
  int prev_cnt = producer_event.stats_cnt;
  while (prev_cnt == producer_event.stats_cnt) {
    Test::Say("Waiting for final producer stats event\n");
    p->poll(100);
  }

  /*
   * Consumer:
   * Wait for a one last stats emit when all messages have been received,
   * since previous stats may have been enqueued but not served we
   * skip the first 2.
   */
  prev_cnt = consumer_event.stats_cnt;
  while (prev_cnt + 2 >= consumer_event.stats_cnt) {
    Test::Say(tostr() << "Waiting for final consumer stats event: " <<
              consumer_event.stats_cnt << "\n");
    c->poll(100);
  }


  verify_e2e_stats(producer_event.last, consumer_event.last,
                   exp_parts, partcnt);


  c->close();

  delete p;
  delete c;
}
#endif

extern "C" {
  int main_0053_stats_timing (int argc, char **argv) {
    test_stats_timing();
    return 0;
  }

  int main_0053_stats (int argc, char **argv) {
#if WITH_RAPIDJSON
    test_stats();
#else
    Test::Skip("RapidJSON >=1.1.0 not available\n");
#endif
    return 0;
  }
}