/*
 * Copyright 2017 MapD Technologies, 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.
 */

/**
 * @file    StreamInsert.cpp
 * @author  Wei Hong <wei@mapd.com>
 * @brief   Sample MapD Client code for inserting a stream of rows
 * with optional transformations from stdin to a MapD table.
 *
 * Copyright (c) 2014 MapD Technologies, Inc.  All rights reserved.
 **/

#include <cstring>
#include <string>
#include <iostream>
#include <iterator>
#include <boost/regex.hpp>

#include <thread>
#include <chrono>

#include <boost/program_options.hpp>

// include files for Thrift and MapD Thrift Services
#include "gen-cpp/MapD.h"
#include <thrift/transport/TSocket.h>
#include <thrift/protocol/TBinaryProtocol.h>
#include <thrift/transport/TBufferTransports.h>

using namespace ::apache::thrift;
using namespace ::apache::thrift::protocol;
using namespace ::apache::thrift::transport;

// Thrift uses boost::shared_ptr instead of std::shared_ptr
using boost::shared_ptr;

struct CopyParams {
  char delimiter;
  std::string null_str;
  char line_delim;
  size_t batch_size;
  size_t retry_count;
  size_t retry_wait;
  CopyParams(char d, const std::string& n, char l, size_t b, size_t retries, size_t wait)
      : delimiter(d), null_str(n), line_delim(l), batch_size(b), retry_count(retries), retry_wait(wait) {}
};

struct ConnectionDetails {
  std::string server_host;
  int port;
  std::string db_name;
  std::string user_name;
  std::string passwd;
  ConnectionDetails(std::string in_server_host,
                    int in_port,
                    std::string in_db_name,
                    std::string in_user_name,
                    std::string in_passwd)
      : server_host(in_server_host), port(in_port), db_name(in_db_name), user_name(in_user_name), passwd(in_passwd) {}
};

bool print_error_data = false;
bool print_transformation = false;

shared_ptr<MapDClient> client;
TSessionId session;
shared_ptr<apache::thrift::transport::TTransport> mytransport;

namespace {
// anonymous namespace for private functions

#define MAX_FIELD_LEN 20000

void createConnection(ConnectionDetails con) {
  shared_ptr<TTransport> socket(new TSocket(con.server_host, con.port));
  mytransport.reset(new TBufferedTransport(socket));
  shared_ptr<TProtocol> protocol(new TBinaryProtocol(mytransport));
  client.reset(new MapDClient(protocol));
  try {
    mytransport->open();                                               // open transport
    client->connect(session, con.user_name, con.passwd, con.db_name);  // connect to mapd_server
  } catch (TMapDException& e) {
    std::cerr << e.error_msg << std::endl;
  } catch (TException& te) {
    std::cerr << "Thrift error: " << te.what() << std::endl;
  }
}

void closeConnection() {
  try {
    client->disconnect(session);  // disconnect from mapd_server
    mytransport->close();         // close transport
  } catch (TMapDException& e) {
    std::cerr << e.error_msg << std::endl;
  } catch (TException& te) {
    std::cerr << "Thrift error: " << te.what() << std::endl;
  }
}

void wait_disconnet_reconnnect_retry(size_t tries, CopyParams copy_params, ConnectionDetails conn_details) {
  std::cout << "  Waiting  " << copy_params.retry_wait << " secs to retry Inserts , will try "
            << (copy_params.retry_count - tries) << " times more " << std::endl;
  sleep(copy_params.retry_wait);

  closeConnection();
  createConnection(conn_details);
}

void do_load(int& nrows,
             int& nskipped,
             std::vector<TStringRow> input_rows,
             const std::string& table_name,
             CopyParams copy_params,
             ConnectionDetails conn_details) {
  for (size_t tries = 0; tries < copy_params.retry_count; tries++) {  // allow for retries in case of insert failure
    try {
      client->load_table(session, table_name, input_rows);
      nrows += input_rows.size();
      std::cout << nrows << " Rows Inserted, " << nskipped << " rows skipped." << std::endl;
      // we successfully loaded the data, lets move on
      return;
    } catch (TMapDException& e) {
      std::cerr << "Exception trying to insert data " << e.error_msg << std::endl;
      wait_disconnet_reconnnect_retry(tries, copy_params, conn_details);
    } catch (TException& te) {
      std::cerr << "Exception trying to insert data " << te.what() << std::endl;
      wait_disconnet_reconnnect_retry(tries, copy_params, conn_details);
    }
  }
  std::cerr << "Retries exhausted program terminated" << std::endl;
  exit(1);
}

// reads copy_params.delimiter delimited rows from std::cin and load them to
// table_name in batches of size copy_params.batch_size until EOF
void stream_insert(const std::string& table_name,
                   const TRowDescriptor& row_desc,
                   const std::map<std::string, std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>>&
                       transformations,
                   const CopyParams& copy_params,
                   const ConnectionDetails conn_details) {
  std::vector<TStringRow> input_rows;
  TStringRow row;

  std::ios_base::sync_with_stdio(false);
  std::istream_iterator<char> eos;
  std::cin >> std::noskipws;
  std::istream_iterator<char> iit(std::cin);

  char field[MAX_FIELD_LEN];
  size_t field_i = 0;

  int nrows = 0;
  int nskipped = 0;

  const std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>* xforms[row_desc.size()];
  for (size_t i = 0; i < row_desc.size(); i++) {
    auto it = transformations.find(row_desc[i].col_name);
    if (it != transformations.end())
      xforms[i] = &(it->second);
    else
      xforms[i] = nullptr;
  }

  while (iit != eos) {
    row.cols.clear();
    // construct a row
    while (iit != eos) {
      if (*iit == copy_params.delimiter || *iit == copy_params.line_delim) {
        bool end_of_field = (*iit == copy_params.delimiter);
        bool end_of_row;
        if (end_of_field)
          end_of_row = false;
        else {
          end_of_row =
              (row_desc[row.cols.size()].col_type.type != TDatumType::STR) || (row.cols.size() == row_desc.size() - 1);
          if (!end_of_row) {
            size_t l = copy_params.null_str.size();
            if (field_i >= l && strncmp(field + field_i - l, copy_params.null_str.c_str(), l) == 0) {
              end_of_row = true;
              // std::cout << "new line after null.\n";
            }
          }
        }
        if (!end_of_field && !end_of_row) {
          // not enough columns yet and it is a string column
          // treat the line delimiter as part of the string
          field[field_i++] = *iit;
        } else {
          field[field_i] = '\0';
          field_i = 0;
          TStringValue ts;
          ts.str_val = std::string(field);
          ts.is_null = (ts.str_val.empty() || ts.str_val == copy_params.null_str);
          auto xform = row.cols.size() < row_desc.size() ? xforms[row.cols.size()] : nullptr;
          if (!ts.is_null && xform != nullptr) {
            if (print_transformation)
              std::cout << "\ntransforming\n" << ts.str_val << "\nto\n";
            ts.str_val = boost::regex_replace(ts.str_val, *xform->first, *xform->second);
            if (ts.str_val.empty())
              ts.is_null = true;
            if (print_transformation)
              std::cout << ts.str_val << std::endl;
          }
          row.cols.push_back(ts);  // add column value to row
          if (end_of_row || (row.cols.size() > row_desc.size()))
            break;  // found row
        }
      } else {
        field[field_i++] = *iit;
      }
      if (field_i >= MAX_FIELD_LEN) {
        field[MAX_FIELD_LEN - 1] = '\0';
        std::cerr << "String too long for buffer." << std::endl;
        if (print_error_data)
          std::cerr << field << std::endl;
        field_i = 0;
        break;
      }
      ++iit;
    }
    if (row.cols.size() == row_desc.size()) {
      input_rows.push_back(row);
      if (input_rows.size() >= copy_params.batch_size) {
        do_load(nrows, nskipped, input_rows, table_name, copy_params, conn_details);
        input_rows.clear();
      }
    } else {
      ++nskipped;
      if (print_error_data) {
        std::cerr << "Incorrect number of columns for row at: ";
        bool not_first = false;
        for (const auto& p : row.cols) {
          if (not_first)
            std::cerr << copy_params.delimiter;
          else
            not_first = true;
          std::cerr << &p;
        }
        std::cerr << std::endl;
      }
      if (row.cols.size() > row_desc.size()) {
        // skip to the next line delimiter
        while (*iit != copy_params.line_delim)
          ++iit;
      }
    }
    ++iit;
  }
  // load remaining rowset if any
  if (input_rows.size() > 0) {
    do_load(nrows, nskipped, input_rows, table_name, copy_params, conn_details);
  }
}
}

int main(int argc, char** argv) {
  std::string server_host("localhost");  // default to localhost
  int port = 9091;                       // default port number
  std::string table_name;
  std::string db_name;
  std::string user_name;
  std::string passwd;
  std::string delim_str(","), nulls("\\N"), line_delim_str("\n");
  size_t batch_size = 10000;
  size_t retry_count = 10;
  size_t retry_wait = 5;
  std::vector<std::string> xforms;
  std::map<std::string, std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>> transformations;

  namespace po = boost::program_options;

  po::options_description desc("Options");
  desc.add_options()("help,h", "Print help messages ");
  desc.add_options()("table", po::value<std::string>(&table_name)->required(), "Table Name");
  desc.add_options()("database", po::value<std::string>(&db_name)->required(), "Database Name");
  desc.add_options()("user,u", po::value<std::string>(&user_name)->required(), "User Name");
  desc.add_options()("passwd,p", po::value<std::string>(&passwd)->required(), "User Password");
  desc.add_options()("host", po::value<std::string>(&server_host)->default_value(server_host), "MapD Server Hostname");
  desc.add_options()("port", po::value<int>(&port)->default_value(port), "MapD Server Port Number");
  desc.add_options()("delim", po::value<std::string>(&delim_str)->default_value(delim_str), "Field delimiter");
  desc.add_options()("null", po::value<std::string>(&nulls), "NULL string");
  desc.add_options()("line", po::value<std::string>(&line_delim_str), "Line delimiter");
  desc.add_options()("batch", po::value<size_t>(&batch_size)->default_value(batch_size), "Insert batch size");
  desc.add_options()(
      "retry_count", po::value<size_t>(&retry_count)->default_value(retry_count), "Number of time to retry an insert");
  desc.add_options()(
      "retry_wait", po::value<size_t>(&retry_wait)->default_value(retry_wait), "wait in secs between retries");
  desc.add_options()(
      "transform,t", po::value<std::vector<std::string>>(&xforms)->multitoken(), "Column Transformations");
  desc.add_options()("print_error", "Print Error Rows");
  desc.add_options()("print_transform", "Print Transformations");

  po::positional_options_description positionalOptions;
  positionalOptions.add("table", 1);
  positionalOptions.add("database", 1);

  po::variables_map vm;

  try {
    po::store(po::command_line_parser(argc, argv).options(desc).positional(positionalOptions).run(), vm);
    if (vm.count("help")) {
      std::cout
          << "Usage: <table name> <database name> {-u|--user} <user> {-p|--passwd} <password> [{--host} "
             "<hostname>][--port <port number>][--delim <delimiter>][--null <null string>][--line <line "
             "delimiter>][--batch <batch size>][{-t|--transform} transformation "
             "...][--retry_count <num_of_retries>] [--retry_wait <wait in secs>][--print_error][--print_transform]\n\n";
      std::cout << desc << std::endl;
      return 0;
    }
    if (vm.count("print_error"))
      print_error_data = true;
    if (vm.count("print_transform"))
      print_transformation = true;

    po::notify(vm);
  } catch (boost::program_options::error& e) {
    std::cerr << "Usage Error: " << e.what() << std::endl;
    return 1;
  }

  char delim = delim_str[0];
  if (delim == '\\') {
    if (delim_str.size() < 2 || (delim_str[1] != 'x' && delim_str[1] != 't' && delim_str[1] != 'n')) {
      std::cerr << "Incorrect delimiter string: " << delim_str << std::endl;
      return 1;
    }
    if (delim_str[1] == 't')
      delim = '\t';
    else if (delim_str[1] == 'n')
      delim = '\n';
    else {
      std::string d(delim_str);
      d[0] = '0';
      delim = (char)std::stoi(d, nullptr, 16);
    }
  }
  if (isprint(delim))
    std::cout << "Field Delimiter: " << delim << std::endl;
  else if (delim == '\t')
    std::cout << "Field Delimiter: "
              << "\\t" << std::endl;
  else if (delim == '\n')
    std::cout << "Field Delimiter: "
              << "\\n"
              << std::endl;
  else
    std::cout << "Field Delimiter: \\x" << std::hex << (int)delim << std::endl;
  char line_delim = line_delim_str[0];
  if (line_delim == '\\') {
    if (line_delim_str.size() < 2 ||
        (line_delim_str[1] != 'x' && line_delim_str[1] != 't' && line_delim_str[1] != 'n')) {
      std::cerr << "Incorrect delimiter string: " << line_delim_str << std::endl;
      return 1;
    }
    if (line_delim_str[1] == 't')
      line_delim = '\t';
    else if (delim_str[1] == 'n')
      line_delim = '\n';
    else {
      std::string d(delim_str);
      d[0] = '0';
      line_delim = (char)std::stoi(d, nullptr, 16);
    }
  }
  if (isprint(line_delim))
    std::cout << "Line Delimiter: " << line_delim << std::endl;
  else if (line_delim == '\t')
    std::cout << "Line Delimiter: "
              << "\\t" << std::endl;
  else if (line_delim == '\n')
    std::cout << "Line Delimiter: "
              << "\\n"
              << std::endl;
  else
    std::cout << "Line Delimiter: \\x" << std::hex << (int)line_delim << std::endl;
  std::cout << "Null String: " << nulls << std::endl;
  std::cout << "Insert Batch Size: " << std::dec << batch_size << std::endl;

  for (auto& t : xforms) {
    auto n = t.find_first_of(':');
    if (n == std::string::npos) {
      std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/" << std::endl;
      return 1;
    }
    std::string col_name = t.substr(0, n);
    if (t.size() < n + 3 || t[n + 1] != 's' || t[n + 2] != '/') {
      std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/" << std::endl;
      return 1;
    }
    auto n1 = n + 3;
    auto n2 = t.find_first_of('/', n1);
    if (n2 == std::string::npos) {
      std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/" << std::endl;
      return 1;
    }
    std::string regex_str = t.substr(n1, n2 - n1);
    n1 = n2 + 1;
    n2 = t.find_first_of('/', n1);
    if (n2 == std::string::npos) {
      std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/" << std::endl;
      return 1;
    }
    std::string fmt_str = t.substr(n1, n2 - n1);
    std::cout << "transform " << col_name << ": s/" << regex_str << "/" << fmt_str << "/" << std::endl;
    transformations[col_name] = std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>(
        std::unique_ptr<boost::regex>(new boost::regex(regex_str)),
        std::unique_ptr<std::string>(new std::string(fmt_str)));
  }

  CopyParams copy_params(delim, nulls, line_delim, batch_size, retry_count, retry_wait);

  // for attaching debugger std::this_thread::sleep_for (std::chrono::seconds(20));
  ConnectionDetails conn_details(server_host, port, db_name, user_name, passwd);
  createConnection(conn_details);

  TTableDetails table_details;
  client->get_table_details(table_details, session, table_name);
  stream_insert(table_name, table_details.row_desc, transformations, copy_params, conn_details);

  closeConnection();

  return 0;
}