#include "main.hpp"

#include <chrono>
#include <fstream>

#include "arguments.hpp"
#include "arguments_parser.hpp"
#include "benchmark.hpp"
#include "driver.hpp"
#include "endpoint.hpp"
#include "endpoint_parser.hpp"
#include "json_parser.hpp"
#include "logger.hpp"
#include "pod_type.hpp"
#include "result.hpp"
#include "test.hpp"
#include "test_manager.hpp"

#include "pods/trtechargableevents.hpp"
#include "pods/trtecounters.hpp"
#include "pods/ttaconsumptionbuckets.hpp"

int main(int argc, const char* argv[]) {
  logger& l = logger::instance();

  arguments args;

  arguments_parser parser(args);
  if (!parser(argc, argv)) {
    l(log_level::error) << "Failed to parse arguments." << std::endl;
    return EXIT_FAILURE;
  }

  test_manager manager;

  if (args.list_tests) {
    std::vector<std::reference_wrapper<test>> tests;
    manager.tests(tests);

    for (const test& t : tests) {
      std::cout << "- " << t << std::endl;
    }

    return EXIT_SUCCESS;
  }

  if (args.databases.empty()) {
    l(log_level::error) << "List of databases is empty." << std::endl;
    return EXIT_FAILURE;
  }

  std::size_t maximum_space = args.databases.size();

  std::vector<driver::client> clients;
  clients.reserve(maximum_space);

  std::vector<stdx::optional<driver::database>> databases;
  databases.reserve(maximum_space);

  std::transform(args.databases.begin(), args.databases.end(),
                 std::back_inserter(databases),
                 [&l, &clients](const stdx::string_view& database)
                     -> stdx::optional<driver::database> {

                       endpoint e;
                       endpoint_parser parser(e);

                       if (!parser(database)) {
                         l(log_level::warning)
                             << "Failed to parse database uri." << std::endl;
                         return {};
                       }

                       driver::type type;
                       try {
                         type = driver::types.at(e.type);
                       } catch (...) {
                         l(log_level::warning) << "Invalid database type ("
                                               << e.type << ")." << std::endl;
                         return {};
                       }

                       clients.emplace_back(type);
                       driver::client& cl = clients.back();

                       auto res = cl.connect(std::string(e.host), e.port);
                       if (!res) {
                         l(log_level::warning)
                             << "Cannot connect to the database." << std::endl;
                         l(log_level::warning)
                             << res.what() << " ("
                             << static_cast<int>(res.get_error_code()) << ")"
                             << std::endl;

                         clients.pop_back();
                         return {};
                       }

                       driver::database db = cl[e.database.data()];
                       if (!db) {
                         l(log_level::warning) << "Missing database (" << e.type
                                               << ")." << std::endl;

                         clients.pop_back();
                         return {};
                       }

                       return db;
                     });

  int available_databases = std::count_if(
      databases.begin(), databases.end(),
      [](const stdx::optional<driver::database>& db) { return db; });

  if (available_databases == 0) {
    l(log_level::error) << "Any database is not available." << std::endl;
    return EXIT_FAILURE;
  }

  std::map<pod_type, std::vector<driver::object>> datasets;

  std::for_each(args.datasets.begin(), args.datasets.end(),
                [&](const stdx::string_view& dataset) {
                  std::fstream file(dataset.data());

                  if (!file) {
                    l(log_level::warning) << "File \"" << dataset
                                          << "\" is missing." << std::endl;
                    return;
                  }

                  std::string header;
                  if (!std::getline(file, header)) {
                    l(log_level::warning) << "Header is missing." << std::endl;
                    return;
                  }

                  pod_type type;
                  try {
                    type = pod_types.at(header);
                  } catch (...) {
                    l(log_level::warning) << "Invalid header type."
                                          << std::endl;
                    return;
                  }

                  json_parser parser;

                  for (std::string line; std::getline(file, line);) {
                    datasets[type].push_back(parser(line));
                  }
                });

  int available_datasets = std::count_if(
      datasets.begin(), datasets.end(),
      [](const std::pair<pod_type, std::vector<driver::object>>& pair) {
        return !pair.second.empty();
      });

  if (available_datasets == 0) {
    l(log_level::error) << "Any dataset is not available." << std::endl;
    return EXIT_FAILURE;
  }

  std::map<driver::type, double> bulk_times;

  std::for_each(
      datasets.begin(), datasets.end(),
      [&](const std::pair<pod_type, std::vector<driver::object>>& pair) {
        for (const auto& db : databases) {
          if (!db) {
            continue;
          }

          stdx::string_view table;
          try {
            table = pod_tables.at(pair.first);
          } catch (...) {
            l(log_level::error) << "Missing table name." << std::endl;
            std::exit(EXIT_FAILURE);
          }

          auto& cl = db->get_client();

          std::chrono::time_point<std::chrono::system_clock> start =
              std::chrono::system_clock::now();

          auto res = cl.bulk_refs(db.value(), std::string(table), pair.second);
          if (!res) {
            l(log_level::warning)
                << "Cannot bulk " << to_string(pair.first) << " to database ("
                << std::string(driver::to_string(cl.get_driver_type())) << ")."
                << std::endl;
          }

          std::chrono::time_point<std::chrono::system_clock> end =
              std::chrono::system_clock::now();

          std::chrono::duration<double> elapsed_seconds = end - start;
          bulk_times[cl.get_driver_type()] = elapsed_seconds.count();
        }
      });

  std::cout << "import-datasets" << std::endl;
  std::for_each(bulk_times.begin(), bulk_times.end(),
                [](const std::pair<driver::type, double>& pair) {
                  std::cout << std::string(driver::to_string(pair.first))
                            << ": " << pair.second << std::endl;
                });
  std::cout << std::endl;

  std::vector<stdx::optional<std::reference_wrapper<test>>> tests;
  tests.reserve(args.tests.size());

  std::transform(args.tests.begin(), args.tests.end(),
                 std::back_inserter(tests),
                 [&manager, &l](const stdx::string_view& test_name)
                     -> stdx::optional<std::reference_wrapper<test>> {
                       test& t = manager[test_name];
                       if (!t) {
                         l(log_level::warning) << "Invalid test." << std::endl;
                         return {};
                       }

                       return {t};
                     });

  int available_tests = std::count_if(
      tests.begin(), tests.end(),
      [](const stdx::optional<std::reference_wrapper<test>>& t) { return t; });

  if (available_tests == 0) {
    l(log_level::error) << "Any test is not available." << std::endl;
    return EXIT_FAILURE;
  }

  std::vector<result> results;
  results.reserve(available_tests);

  std::for_each(tests.begin(), tests.end(),
                [&](const stdx::optional<std::reference_wrapper<test>>& opt) {
                  if (!opt) {
                    return;
                  }

                  test& t = opt.value();
                  benchmark bench(t);

                  result res(t, args);

                  for (const auto& db : databases) {
                    if (!db) {
                      continue;
                    }

                    auto& cl = db->get_client();

                    for (std::size_t i = 0; i < args.iterations; i++) {
                      res.add_time(cl.get_driver_type(), bench(db.value()));
                    }
                  }

                  results.emplace_back(res);
                });

  std::cout << std::fixed;
  for (const auto& res : results) {
    std::cout << res << std::endl;
  }

  return EXIT_SUCCESS;
}