/*
 * Copyright 2014 MongoDB, Inc.
 *
 * 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.
 */

/*
 * This program will scan each BSON document contained in the provided files
 * and print metrics to STDOUT.
 */

#include <bson.h>
#include <stdio.h>
#include <math.h>

#define MAX_RECURSION 100

static double
dtimeofday (void)
{
    struct timeval timeval;
    bson_gettimeofday(&timeval);
    return (timeval.tv_sec + timeval.tv_usec * 0.000001);
}

typedef struct
{
  uint64_t count;
  char *description;
} bson_type_metrics_t;

typedef struct
{
   uint64_t doc_count;
   uint64_t element_count;
   uint64_t doc_size_max;
   uint64_t key_size_tally;
   uint64_t utf8_size_tally;
   uint32_t depth;
   bson_type_metrics_t bson_type_metrics[256];
} bson_metrics_state_t;

static bson_metrics_state_t initial_state = {
   0L, 0L, 0L, 0L, 0L, 0L,
   {
       { /* BSON_TYPE_EOD        = 0x00 */ 0L, "End of document" },
       { /* BSON_TYPE_DOUBLE     = 0x01 */ 0L, "Floating point" },
       { /* BSON_TYPE_UTF8       = 0x02 */ 0L, "UTF-8 string" },
       { /* BSON_TYPE_DOCUMENT   = 0x03 */ 0L, "Embedded document" },
       { /* BSON_TYPE_ARRAY      = 0x04 */ 0L, "Array" },
       { /* BSON_TYPE_BINARY     = 0x05 */ 0L, "Binary data" },
       { /* BSON_TYPE_UNDEFINED  = 0x06 */ 0L, "Undefined - Deprecated" },
       { /* BSON_TYPE_OID        = 0x07 */ 0L, "ObjectId" },
       { /* BSON_TYPE_BOOL       = 0x08 */ 0L, "Boolean" },
       { /* BSON_TYPE_DATE_TIME  = 0x09 */ 0L, "UTC datetime" },
       { /* BSON_TYPE_NULL       = 0x0A */ 0L, "Null value" },
       { /* BSON_TYPE_REGEX      = 0x0B */ 0L, "Regular expression" },
       { /* BSON_TYPE_DBPOINTER  = 0x0C */ 0L, "DBPointer - Deprecated" },
       { /* BSON_TYPE_CODE       = 0x0D */ 0L, "JavaScript code" },
       { /* BSON_TYPE_SYMBOL     = 0x0E */ 0L, "Symbol - Deprecated" },
       { /* BSON_TYPE_CODEWSCOPE = 0x0F */ 0L, "JavaScript code w/ scope" },
       { /* BSON_TYPE_INT32      = 0x10 */ 0L, "32-bit Integer" },
       { /* BSON_TYPE_TIMESTAMP  = 0x11 */ 0L, "Timestamp" },
       { /* BSON_TYPE_INT64      = 0x12 */ 0L, "64-bit Integer" },
       { 0L, NULL }
   }
};

static bson_metrics_state_t state;

static int
compar_bson_type_metrics (const void *a,
                          const void *b)
{
    return (((bson_type_metrics_t*)b)->count - ((bson_type_metrics_t*)a)->count);
}

/*
 * Forward declarations.
 */
static bool bson_metrics_visit_array (const bson_iter_t *iter,
                                      const char        *key,
                                      const bson_t      *v_array,
                                      void              *data);
static bool bson_metrics_visit_document (const bson_iter_t *iter,
                                         const char        *key,
                                         const bson_t      *v_document,
                                         void              *data);

static bool
bson_metrics_visit_utf8 (const bson_iter_t *iter,
                         const char        *key,
                         size_t             v_utf8_len,
                         const char        *v_utf8,
                         void              *data)
{
   bson_metrics_state_t *state = data;
   state->utf8_size_tally += v_utf8_len;

   return false;
}

static bool
bson_metrics_visit_before (const bson_iter_t *iter,
                           const char        *key,
                           void              *data)
{
   bson_metrics_state_t *state = data;
   bson_type_t btype;
   ++state->element_count;
   state->key_size_tally += strlen(key); /* TODO - filter out array keys(?) */
   btype = bson_iter_type (iter);
   ++state->bson_type_metrics[btype].count;

   return false;
}

static const bson_visitor_t bson_metrics_visitors = {
   bson_metrics_visit_before,
   NULL, /* visit_after */
   NULL, /* visit_corrupt */
   NULL, /* visit_double */
   bson_metrics_visit_utf8,
   bson_metrics_visit_document,
   bson_metrics_visit_array,
   NULL, /* visit_binary */
   NULL, /* visit_undefined */
   NULL, /* visit_oid */
   NULL, /* visit_bool */
   NULL, /* visit_date_time */
   NULL, /* visit_null */
   NULL, /* visit_regex */
   NULL, /* visit_dbpointer */
   NULL, /* visit_code */
   NULL, /* visit_symbol */
   NULL, /* visit_codewscope */
   NULL, /* visit_int32 */
   NULL, /* visit_timestamp */
   NULL, /* visit_int64 */
   NULL, /* visit_maxkey */
   NULL, /* visit_minkey */
};

static bool
bson_metrics_visit_document (const bson_iter_t *iter,
                             const char        *key,
                             const bson_t      *v_document,
                             void              *data)
{
   bson_metrics_state_t *state = data;
   bson_iter_t child;

   if (state->depth >= MAX_RECURSION) {
      fprintf (stderr, "Invalid document, max recursion reached.\n");
      return true;
   }

   if (bson_iter_init (&child, v_document)) {
      state->depth++;
      bson_iter_visit_all (&child, &bson_metrics_visitors, data);
      state->depth--;
   }

   return false;
}

static bool
bson_metrics_visit_array (const bson_iter_t *iter,
                          const char        *key,
                          const bson_t      *v_array,
                          void              *data)
{
   bson_metrics_state_t *state = data;
   bson_iter_t child;

   if (state->depth >= MAX_RECURSION) {
      fprintf (stderr, "Invalid document, max recursion reached.\n");
      return true;
   }

   if (bson_iter_init (&child, v_array)) {
      state->depth++;
      bson_iter_visit_all (&child, &bson_metrics_visitors, data);
      state->depth--;
   }

   return false;
}

static void
bson_metrics (const bson_t *bson,
              size_t       *length,
              void         *data)
{
   bson_metrics_state_t *state = data;
   ++state->doc_count;
   bson_iter_t iter;

   if (bson_iter_init (&iter, bson)) {
      bson_iter_visit_all (&iter, &bson_metrics_visitors, data);
   }
}

int
main (int   argc,
      char *argv[])
{
   bson_reader_t *reader;
   const bson_t *b;
   bson_error_t error;
   const char *filename;
   int i, j;
   double dtime_before, dtime_after, dtime_delta;
   uint64_t aggregate_count;
   off_t mark;

   /*
    * Print program usage if no arguments are provided.
    */
   if (argc == 1) {
      fprintf(stderr, "usage: %s FILE...\n", argv[0]);
      return 1;
   }

   /*
    * Process command line arguments expecting each to be a filename.
    */
   printf("[");
   for (i = 1; i < argc; i++) {
      if (i > 1) printf(",");
      filename = argv[i];

      /*
       * Initialize a new reader for this file descriptor.
       */
      if (!(reader = bson_reader_new_from_file (filename, &error))) {
         fprintf (stderr, "Failed to open \"%s\": %s\n",
                  filename, error.message);
         continue;
      }

      state = initial_state;
      dtime_before = dtimeofday();
      mark = 0;
      while ((b = bson_reader_read (reader, NULL))) {
        off_t pos = bson_reader_tell(reader);
        state.doc_size_max = BSON_MAX(pos - mark, state.doc_size_max);
        mark = pos;
        bson_metrics(b, NULL, &state);
      }
      dtime_after = dtimeofday();
      dtime_delta = BSON_MAX(dtime_after - dtime_before, 0.000001);
      state.bson_type_metrics[BSON_TYPE_MAXKEY].description = "Max key";
      state.bson_type_metrics[BSON_TYPE_MINKEY].description = "Min key";
      aggregate_count = state.bson_type_metrics[BSON_TYPE_DOCUMENT].count + state.bson_type_metrics[BSON_TYPE_ARRAY].count;
      qsort(state.bson_type_metrics, 256, sizeof(bson_type_metrics_t), compar_bson_type_metrics);

      printf("\n  {\n");
      printf("    \"file\": \"%s\",\n", filename);
      printf("    \"secs\": %.2f,\n", dtime_delta);
      printf("    \"docs_per_sec\": %"PRIu64",\n", (uint64_t)round(state.doc_count/dtime_delta));
      printf("    \"docs\": %"PRIu64",\n", state.doc_count);
      printf("    \"elements\": %"PRIu64",\n", state.element_count);
      printf("    \"elements_per_doc\": %"PRIu64",\n", (uint64_t)round((double)state.element_count/(double)BSON_MAX(state.doc_count, 1)));
      printf("    \"aggregates\": %"PRIu64",\n", aggregate_count);
      printf("    \"aggregates_per_doc\": %"PRIu64",\n", (uint64_t)round((double)aggregate_count/(double)BSON_MAX(state.doc_count, 1)));
      printf("    \"degree\": %"PRIu64",\n", (uint64_t)round((double)state.element_count/((double)BSON_MAX(state.doc_count + aggregate_count, 1))));
      printf("    \"doc_size_max\": %"PRIu64",\n", state.doc_size_max);
      printf("    \"doc_size_average\": %"PRIu64",\n", (uint64_t)round((double)bson_reader_tell(reader)/(double)BSON_MAX(state.doc_count, 1)));
      printf("    \"key_size_average\": %"PRIu64",\n", (uint64_t)round((double)state.key_size_tally/(double)BSON_MAX(state.element_count, 1)));
      printf("    \"string_size_average\": %"PRIu64",\n", (uint64_t)round((double)state.utf8_size_tally/(double)BSON_MAX(state.bson_type_metrics[BSON_TYPE_UTF8].count, 1)));
      printf("    \"percent_by_type\": {\n");
      for (j = 0; state.bson_type_metrics[j].count > 0; j++) {
        bson_type_metrics_t bson_type_metrics = state.bson_type_metrics[j];
        printf("      \"%s\": %"PRIu64",\n", bson_type_metrics.description, (uint64_t)round((double)bson_type_metrics.count*100.0/(double)BSON_MAX(state.element_count, 1)));
      }
      printf("    }\n");
      printf("  }");

      /*
       * Cleanup after our reader, which closes the file descriptor.
       */
      bson_reader_destroy (reader);
   }
   printf("\n]\n");

   return 0;
}
