/**
 * 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 "query_execution/Shiftboss.hpp"

#include <cstddef>
#include <cstdlib>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "catalog/CatalogTypedefs.hpp"
#include "query_execution/QueryContext.hpp"
#include "query_execution/QueryExecutionMessages.pb.h"
#include "query_execution/QueryExecutionTypedefs.hpp"
#include "query_execution/QueryExecutionUtil.hpp"
#include "query_execution/WorkerMessage.hpp"
#include "relational_operators/RebuildWorkOrder.hpp"
#include "relational_operators/WorkOrderFactory.hpp"
#include "storage/InsertDestination.hpp"
#include "storage/StorageBlock.hpp"
#include "storage/StorageManager.hpp"
#include "threading/ThreadUtil.hpp"

#include "glog/logging.h"

#include "tmb/address.h"
#include "tmb/id_typedefs.h"
#include "tmb/message_bus.h"
#include "tmb/message_style.h"
#include "tmb/tagged_message.h"

using std::free;
using std::malloc;
using std::move;
using std::size_t;
using std::string;
using std::unique_ptr;
using std::vector;

using tmb::MessageBus;
using tmb::TaggedMessage;

namespace quickstep {

class WorkOrder;

void Shiftboss::run() {
  if (cpu_id_ >= 0) {
    // We can pin the shiftboss thread to a CPU if specified.
    ThreadUtil::BindToCPU(cpu_id_);
  }

  for (;;) {
    // Receive() is a blocking call, causing this thread to sleep until next
    // message is received.
    AnnotatedMessage annotated_message(bus_->Receive(shiftboss_client_id_, 0, true));
    DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
               << "') received the typed '" << annotated_message.tagged_message.message_type()
               << "' message from client " << annotated_message.sender;
    switch (annotated_message.tagged_message.message_type()) {
      case kShiftbossRegistrationResponseMessage: {
        foreman_client_id_ = annotated_message.sender;

        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::ShiftbossRegistrationResponseMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        shiftboss_index_ = proto.shiftboss_index();
        break;
      }
      case kQueryInitiateMessage: {
        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::QueryInitiateMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        processQueryInitiateMessage(proto.query_id(), proto.catalog_database_cache(), proto.query_context());
        break;
      }
      case kWorkOrderMessage: {
        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::WorkOrderMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        const std::size_t query_id = proto.query_id();
        DCHECK_EQ(1u, query_contexts_.count(query_id));

        WorkOrder *work_order = WorkOrderFactory::ReconstructFromProto(proto.work_order(),
                                                                       &database_cache_,
                                                                       query_contexts_[query_id].get(),
                                                                       storage_manager_,
                                                                       shiftboss_client_id_,
                                                                       bus_);

        unique_ptr<WorkerMessage> worker_message(
            WorkerMessage::WorkOrderMessage(work_order, proto.operator_index()));

        TaggedMessage worker_tagged_message(worker_message.get(),
                                            sizeof(*worker_message),
                                            kWorkOrderMessage);

        const size_t worker_index = getSchedulableWorker();
        DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
                   << "') forwarded WorkOrderMessage (typed '" << kWorkOrderMessage
                   << "') from Foreman to worker " << worker_index;

        const MessageBus::SendStatus send_status =
            QueryExecutionUtil::SendTMBMessage(bus_,
                                               shiftboss_client_id_,
                                               workers_->getClientID(worker_index),
                                               move(worker_tagged_message));
        CHECK(send_status == MessageBus::SendStatus::kOK);
        break;
      }
      case kInitiateRebuildMessage: {
        // Construct rebuild work orders, and send back their number to
        // 'ForemanDistributed'.
        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::InitiateRebuildMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        processInitiateRebuildMessage(proto.query_id(),
                                      proto.operator_index(),
                                      proto.insert_destination_index(),
                                      proto.relation_id());
        break;
      }
      case kCatalogRelationNewBlockMessage:  // Fall through.
      case kDataPipelineMessage:
      case kWorkOrderFeedbackMessage:
      case kWorkOrderCompleteMessage:
      case kRebuildWorkOrderCompleteMessage: {
        DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
                   << "') forwarded typed '" << annotated_message.tagged_message.message_type()
                   << "' message from Worker with TMB client ID '" << annotated_message.sender
                   << "' to Foreman with TMB client ID " << foreman_client_id_;

        DCHECK_NE(foreman_client_id_, tmb::kClientIdNone);
        const MessageBus::SendStatus send_status =
            QueryExecutionUtil::SendTMBMessage(bus_,
                                               shiftboss_client_id_,
                                               foreman_client_id_,
                                               move(annotated_message.tagged_message));
        CHECK(send_status == MessageBus::SendStatus::kOK);
        break;
      }
      case kQueryTeardownMessage: {
        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::QueryTeardownMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        query_contexts_.erase(proto.query_id());
        break;
      }
      case kSaveQueryResultMessage: {
        const TaggedMessage &tagged_message = annotated_message.tagged_message;

        serialization::SaveQueryResultMessage proto;
        CHECK(proto.ParseFromArray(tagged_message.message(), tagged_message.message_bytes()));

        for (int i = 0; i < proto.blocks_size(); ++i) {
          storage_manager_->saveBlockOrBlob(proto.blocks(i));
        }

        // Clean up query execution states, i.e., QueryContext.
        query_contexts_.erase(proto.query_id());

        serialization::SaveQueryResultResponseMessage proto_response;
        proto_response.set_relation_id(proto.relation_id());
        proto_response.set_cli_id(proto.cli_id());

        const size_t proto_response_length = proto_response.ByteSize();
        char *proto_response_bytes = static_cast<char*>(malloc(proto_response_length));
        CHECK(proto_response.SerializeToArray(proto_response_bytes, proto_response_length));

        TaggedMessage message_response(static_cast<const void*>(proto_response_bytes),
                                       proto_response_length,
                                       kSaveQueryResultResponseMessage);
        free(proto_response_bytes);

        DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
                   << "') sent SaveQueryResultResponseMessage (typed '" << kSaveQueryResultResponseMessage
                   << "') to Foreman with TMB client ID " << foreman_client_id_;
        const MessageBus::SendStatus send_status =
            QueryExecutionUtil::SendTMBMessage(bus_,
                                               shiftboss_client_id_,
                                               foreman_client_id_,
                                               move(message_response));
        CHECK(send_status == MessageBus::SendStatus::kOK);
        break;
      }
      case kPoisonMessage: {
        DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
                   << "') forwarded PoisonMessage (typed '" << kPoisonMessage
                   << "') from Foreman to all workers";

        tmb::MessageStyle broadcast_style;
        broadcast_style.Broadcast(true);

        const MessageBus::SendStatus send_status =
            bus_->Send(shiftboss_client_id_,
                       worker_addresses_,
                       broadcast_style,
                       move(annotated_message.tagged_message));
        CHECK(send_status == MessageBus::SendStatus::kOK);
        return;
      }
      default: {
        LOG(FATAL) << "Unknown TMB message type";
      }
    }
  }
}

size_t Shiftboss::getSchedulableWorker() {
  const size_t num_workers = workers_->getNumWorkers();

  size_t curr_worker = start_worker_index_;
  for (;;) {
    if (workers_->getNumQueuedWorkOrders(curr_worker) < max_msgs_per_worker_) {
      start_worker_index_ = (curr_worker + 1) % num_workers;
      // TODO(zuyu): workers_->incrementNumQueuedWorkOrders(curr_worker);
      // But we need a WorkOrder queue first.
      return curr_worker;
    }

    curr_worker = (curr_worker + 1) % num_workers;
  }
}

void Shiftboss::registerWithForeman() {
  tmb::Address all_addresses;
  all_addresses.All(true);

  tmb::MessageStyle style;

  serialization::ShiftbossRegistrationMessage proto;
  proto.set_work_order_capacity(getWorkOrderCapacity());

  const size_t proto_length = proto.ByteSize();
  char *proto_bytes = static_cast<char*>(malloc(proto_length));
  CHECK(proto.SerializeToArray(proto_bytes, proto_length));

  TaggedMessage message(static_cast<const void*>(proto_bytes),
                        proto_length,
                        kShiftbossRegistrationMessage);
  free(proto_bytes);

  DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
             << "') sent ShiftbossRegistrationMessage (typed '" << kShiftbossRegistrationMessage
             << "') to all";
  tmb::MessageBus::SendStatus send_status =
      bus_->Send(shiftboss_client_id_, all_addresses, style, move(message));
  DCHECK(send_status == tmb::MessageBus::SendStatus::kOK);
}

void Shiftboss::processQueryInitiateMessage(
    const std::size_t query_id,
    const serialization::CatalogDatabase &catalog_database_cache_proto,
    const serialization::QueryContext &query_context_proto) {
  database_cache_.update(catalog_database_cache_proto);

  unique_ptr<QueryContext> query_context(
      new QueryContext(query_context_proto,
                       database_cache_,
                       storage_manager_,
                       shiftboss_client_id_,
                       bus_));
  query_contexts_.emplace(query_id, move(query_context));

  serialization::QueryInitiateResponseMessage proto;
  proto.set_query_id(query_id);

  const size_t proto_length = proto.ByteSize();
  char *proto_bytes = static_cast<char*>(malloc(proto_length));
  CHECK(proto.SerializeToArray(proto_bytes, proto_length));

  TaggedMessage message_response(static_cast<const void*>(proto_bytes),
                                 proto_length,
                                 kQueryInitiateResponseMessage);
  free(proto_bytes);

  DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
             << "') sent QueryInitiateResponseMessage (typed '" << kQueryInitiateResponseMessage
             << "') to Foreman with TMB client ID " << foreman_client_id_;
  const MessageBus::SendStatus send_status =
      QueryExecutionUtil::SendTMBMessage(bus_,
                                         shiftboss_client_id_,
                                         foreman_client_id_,
                                         move(message_response));
  CHECK(send_status == MessageBus::SendStatus::kOK);
}

void Shiftboss::processInitiateRebuildMessage(const std::size_t query_id,
                                              const std::size_t op_index,
                                              const QueryContext::insert_destination_id dest_index,
                                              const relation_id rel_id) {
  DCHECK_NE(foreman_client_id_, tmb::kClientIdNone);

  DCHECK_EQ(1u, query_contexts_.count(query_id));
  InsertDestination *insert_destination = query_contexts_[query_id]->getInsertDestination(dest_index);
  DCHECK(insert_destination != nullptr);

  vector<MutableBlockReference> partially_filled_block_refs;
  insert_destination->getPartiallyFilledBlocks(&partially_filled_block_refs);

  serialization::InitiateRebuildResponseMessage proto;
  proto.set_query_id(query_id);
  proto.set_operator_index(op_index);
  proto.set_num_rebuild_work_orders(partially_filled_block_refs.size());
  proto.set_shiftboss_index(shiftboss_index_);

  const size_t proto_length = proto.ByteSize();
  char *proto_bytes = static_cast<char*>(malloc(proto_length));
  CHECK(proto.SerializeToArray(proto_bytes, proto_length));

  TaggedMessage message_response(static_cast<const void*>(proto_bytes),
                                 proto_length,
                                 kInitiateRebuildResponseMessage);
  free(proto_bytes);

  DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
             << "') sent InitiateRebuildResponseMessage (typed '" << kInitiateRebuildResponseMessage
             << "') to Foreman with TMB client ID " << foreman_client_id_;
  const MessageBus::SendStatus send_status =
      QueryExecutionUtil::SendTMBMessage(bus_,
                                         shiftboss_client_id_,
                                         foreman_client_id_,
                                         move(message_response));
  CHECK(send_status == MessageBus::SendStatus::kOK);

  for (size_t i = 0; i < partially_filled_block_refs.size(); ++i) {
    // NOTE(zuyu): Worker releases the memory after the execution of
    // RebuildWorkOrder on the Worker.
    WorkOrder *rebuild_work_order =
        new RebuildWorkOrder(query_id,
                             move(partially_filled_block_refs[i]),
                             op_index,
                             rel_id,
                             shiftboss_client_id_,
                             bus_);

    unique_ptr<WorkerMessage> worker_message(
        WorkerMessage::RebuildWorkOrderMessage(rebuild_work_order, op_index));

    TaggedMessage worker_tagged_message(worker_message.get(),
                                        sizeof(*worker_message),
                                        kRebuildWorkOrderMessage);

    const size_t worker_index = getSchedulableWorker();
    DLOG(INFO) << "Shiftboss (id '" << shiftboss_client_id_
               << "') sent RebuildWorkOrderMessage (typed '" << kRebuildWorkOrderMessage
               << "') to worker " << worker_index;

    const MessageBus::SendStatus send_status =
        QueryExecutionUtil::SendTMBMessage(bus_,
                                           shiftboss_client_id_,
                                           workers_->getClientID(worker_index),
                                           move(worker_tagged_message));
    CHECK(send_status == MessageBus::SendStatus::kOK);
  }
}

}  // namespace quickstep