// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

#include <chrono>  // NOLINT(build/c++11)
#include <iostream>
#include <memory>
#include <string>
#include <thread>  // NOLINT(build/c++11)
#include <utility>
#include <vector>

#include "gflags/gflags.h"
#include "tmb/address.h"
#include "tmb/id_typedefs.h"
#include "tmb/message_bus.h"
#include "tmb/message_style.h"
#include "tmb/priority.h"
#include "tmb/tagged_message.h"
#include "tmbbench/affinity.h"
#include "tmbbench/bus_setup.h"
#include "tmbbench/messages.h"
#include "tmbbench/receiver_thread.h"
#include "tmbbench/sender_thread.h"

// Command-line flags.
DEFINE_uint64(test_duration, 20,
              "Test duration in seconds. May run for longer if receiver "
              "threads have trouble keeping up with send throughput.");
DEFINE_uint64(sender_threads, 1,
              "Number of sender threads to use. Must be at least 1.");
DEFINE_uint64(receiver_threads, 1,
              "Number of receiver threads to use. Must be at least 1.");
DEFINE_string(sender_thread_affinities, "",
              "Optional comma-seperated list of CPU core IDs to affinitize "
              "sender threads to. -1 indicates no affinity.");
DEFINE_string(receiver_thread_affinities, "",
              "Optional comma-seperated list of CPU core IDs to affinitize "
              "receiver threads to. -1 indicates no affinity.");
DEFINE_string(global_affinity_set, "",
              "As an alternative to sender_thread_affinities and "
              "receiver_thread_affinities, this option can be used to set the "
              "affinity of all threads to the specified comma-seperated list "
              "of CPU IDs.");
DEFINE_uint64(message_bytes, 8,
              "Size of messages exchanged via the TMB in bytes. Must be at "
              "least 8.");
DEFINE_bool(delete_messages_immediately, false,
            "Whether to delete messages immediately as they are received, or "
            "to delete separately afterwards.");

int main(int argc, char *argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  if (argc != 1) {
    std::cerr << "Unrecognized command-line arguments.\n";
    return 1;
  }

  if (FLAGS_message_bytes < 8) {
    std::cerr << "message_bytes must be at least 8.\n";
    return 1;
  }

  std::unique_ptr<tmb::MessageBus> message_bus(
      tmbbench::SetupBusAllInOne(FLAGS_delete_messages_immediately));
  if (!message_bus) {
    return 1;
  }

  std::vector<int> global_affinity;
  if (!FLAGS_global_affinity_set.empty()) {
    if (!(FLAGS_sender_thread_affinities.empty()
          && FLAGS_receiver_thread_affinities.empty())) {
      std::cerr << "Can not use global_affinity_set together with "
                   "sender_thread_affinities or receiver_thread_affinities.";
      return 1;
    }
    if (!tmbbench::ParseGlobalAffinityString(FLAGS_global_affinity_set,
                                             &global_affinity)) {
      std::cerr << "Could not parse global_affinity_set.\n";
      return 1;
    }
  }

  std::vector<std::vector<int>> sender_affinities;
  if (FLAGS_sender_threads < 1) {
    std::cerr << "sender_threads must be at least 1.\n";
    return 1;
  }
  if (FLAGS_sender_thread_affinities.empty()) {
    sender_affinities.resize(FLAGS_sender_threads, global_affinity);
  } else {
    if (!tmbbench::ParseAffinityString(FLAGS_sender_thread_affinities,
                                       &sender_affinities)) {
      std::cerr << "Could not parse sender_thread_affinities.\n";
      return 1;
    }
    if (sender_affinities.size() != FLAGS_sender_threads) {
      std::cerr << "sender_thread_affinities must have the same number of "
                   "entries as the number of sender_threads.\n";
      return 1;
    }
  }

  std::vector<std::vector<int>> receiver_affinities;
  if (FLAGS_receiver_threads < 1) {
    std::cerr << "receiver_threads must be at least 1.\n";
    return 1;
  }
  if (FLAGS_receiver_thread_affinities.empty()) {
    receiver_affinities.resize(FLAGS_receiver_threads, global_affinity);
  } else {
    if (!tmbbench::ParseAffinityString(FLAGS_receiver_thread_affinities,
                                       &receiver_affinities)) {
      std::cerr << "Could not parse receiver_thread_affinities.\n";
      return 1;
    }
    if (receiver_affinities.size() != FLAGS_receiver_threads) {
      std::cerr << "receiver_thread_affinities must have the same number of "
                   "entries as the number of receiver_threads.\n";
      return 1;
    }
  }

  tmb::client_id control_thread_id = message_bus->Connect();
  message_bus->RegisterClientAsSender(control_thread_id, 1);

  std::vector<std::unique_ptr<tmbbench::SenderThread>> sender_threads;
  std::vector<std::unique_ptr<tmbbench::ReceiverThreadBase>> receiver_threads;

  for (const std::vector<int> &affinity : receiver_affinities) {
    if (FLAGS_delete_messages_immediately) {
      receiver_threads.emplace_back(
          new tmbbench::ReceiverThread<true>(message_bus.get(), affinity));
    } else {
      receiver_threads.emplace_back(
          new tmbbench::ReceiverThread<false>(message_bus.get(), affinity));
    }
  }
  for (const std::vector<int> &affinity : sender_affinities) {
    sender_threads.emplace_back(new tmbbench::SenderThread(message_bus.get(),
                                                           FLAGS_message_bytes,
                                                           affinity));
    for (const std::unique_ptr<tmbbench::ReceiverThreadBase> &receiver_thread
         : receiver_threads) {
      sender_threads.back()->AddReceiverID(receiver_thread->GetClientID());
    }
  }

  for (const std::unique_ptr<tmbbench::SenderThread> &sender_thread
       : sender_threads) {
    sender_thread->Start();
  }
  for (const std::unique_ptr<tmbbench::ReceiverThreadBase> &receiver_thread
       : receiver_threads) {
    receiver_thread->Start();
  }

  std::this_thread::sleep_for(std::chrono::seconds(FLAGS_test_duration));

  tmb::MessageStyle style;
  for (const std::unique_ptr<tmbbench::SenderThread> &sender_thread
       : sender_threads) {
    tmb::Address address;
    address.AddRecipient(sender_thread->GetClientID());
    tmb::TaggedMessage message(&tmbbench::kPoisonMessage,
                               sizeof(tmbbench::kPoisonMessage),
                               1);
    message_bus->Send(control_thread_id, address, style, std::move(message));
  }
  for (const std::unique_ptr<tmbbench::ReceiverThreadBase> &receiver_thread
       : receiver_threads) {
    tmb::Address address;
    address.AddRecipient(receiver_thread->GetClientID());
    tmb::TaggedMessage message(&tmbbench::kPoisonMessage,
                               sizeof(tmbbench::kPoisonMessage),
                               1);
    message_bus->Send(control_thread_id,
                      address,
                      style,
                      std::move(message),
                      tmb::kDefaultPriority + 1);
  }

  for (const std::unique_ptr<tmbbench::SenderThread> &sender_thread
       : sender_threads) {
    sender_thread->Join();
  }
  for (const std::unique_ptr<tmbbench::ReceiverThreadBase> &receiver_thread
       : receiver_threads) {
    receiver_thread->Join();
  }

  double total_send_throughput = 0;
  for (const std::unique_ptr<tmbbench::SenderThread> &sender_thread
       : sender_threads) {
    total_send_throughput += sender_thread->GetThroughput();
  }
  double total_receive_throughput = 0;
  for (const std::unique_ptr<tmbbench::ReceiverThreadBase> &receiver_thread
       : receiver_threads) {
    total_receive_throughput += receiver_thread->GetThroughput();
  }

  std::cout << "Send throughput:    " << total_send_throughput
            << " messages/s\n";
  std::cout << "Receive throughput: " << total_receive_throughput
            << " messages/s\n";

  return 0;
}