// Copyright 2020 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "cast/streaming/public/session_messenger.h"

#include <algorithm>
#include <chrono>
#include <string>

#include "cast/common/public/message_port.h"
#include "cast/streaming/message_fields.h"
#include "platform/base/trivial_clock_traits.h"
#include "util/json/json_helpers.h"
#include "util/json/json_serialization.h"
#include "util/osp_logging.h"
#include "util/string_util.h"

namespace openscreen::cast {

namespace {

// Default timeout to receive a reply message in response to a request message
// sent by us.
constexpr std::chrono::milliseconds kReplyTimeout(4000);

// Special character indicating message was sent to all receivers or senders.
constexpr char kAnyDestination[] = "*";

void ReplyIfTimedOut(
    int sequence_number,
    std::vector<std::pair<int, SenderSessionMessenger::ReplyCallback>>*
        replies) {
  for (auto it = replies->begin(); it != replies->end(); ++it) {
    if (it->first == sequence_number) {
      OSP_VLOG << "Reply was an error with due to timeout for sequence number: "
               << sequence_number;

      // We erase before handling the callback, since it may invalidate the
      // replies vector.
      SenderSessionMessenger::ReplyCallback callback = std::move(it->second);
      replies->erase(it);
      callback(Error(Error::Code::kMessageTimeout,
                     string_util::StrCat({"message timed out; max delay of ",
                                          ToString(kReplyTimeout)})));
      return;
    }
  }
}

}  // namespace

SessionMessenger::SessionMessenger(MessagePort& message_port,
                                   std::string source_id,
                                   ErrorCallback cb)
    : message_port_(message_port),
      source_id_(source_id),
      error_callback_(std::move(cb)) {
  OSP_CHECK(!source_id_.empty());
  message_port_->SetClient(*this);
}

SessionMessenger::~SessionMessenger() {
  message_port_->ResetClient();
}

Error SessionMessenger::SendMessage(const std::string& destination_id,
                                    const std::string& namespace_,
                                    const Json::Value& message_root) {
  OSP_CHECK(namespace_ == kCastRemotingNamespace ||
            namespace_ == kCastWebrtcNamespace);
  auto body_or_error = json::Stringify(message_root);
  if (body_or_error.is_error()) {
    return std::move(body_or_error.error());
  }
  OSP_VLOG << "Sending message: DESTINATION[" << destination_id
           << "], NAMESPACE[" << namespace_ << "], BODY:\n"
           << body_or_error.value();
  message_port_->PostMessage(destination_id, namespace_, body_or_error.value());
  return Error::None();
}

void SessionMessenger::ReportError(const Error& error) {
  error_callback_(error);
}

SenderSessionMessenger::SenderSessionMessenger(MessagePort& message_port,
                                               std::string source_id,
                                               std::string receiver_id,
                                               ErrorCallback cb,
                                               TaskRunner& task_runner)
    : SessionMessenger(message_port, std::move(source_id), std::move(cb)),
      task_runner_(task_runner),
      receiver_id_(std::move(receiver_id)) {}

void SenderSessionMessenger::SetHandler(ReceiverMessage::Type type,
                                        ReplyCallback cb) {
  // Currently the only handlers allowed are for RPC and INPUT messages.
  if (type == ReceiverMessage::Type::kRpc) {
    rpc_callback_ = std::move(cb);
  } else if (type == ReceiverMessage::Type::kInput) {
    input_callback_ = std::move(cb);
  } else {
    OSP_NOTREACHED();
  }
}

void SenderSessionMessenger::ResetHandler(ReceiverMessage::Type type) {
  if (type == ReceiverMessage::Type::kRpc) {
    rpc_callback_ = {};
  } else if (type == ReceiverMessage::Type::kInput) {
    input_callback_ = {};
  } else {
    OSP_NOTREACHED();
  }
}

Error SenderSessionMessenger::SendOutboundMessage(SenderMessage message) {
  const auto namespace_ = (message.type == SenderMessage::Type::kRpc ||
                           message.type == SenderMessage::Type::kInput)
                              ? kCastRemotingNamespace
                              : kCastWebrtcNamespace;

  ErrorOr<Json::Value> jsonified = message.ToJson();
  OSP_CHECK(jsonified.is_value()) << "Tried to send an invalid message";
  return SessionMessenger::SendMessage(receiver_id_, namespace_,
                                       jsonified.value());
}

Error SenderSessionMessenger::SendRpcMessage(ByteView message) {
  return SendOutboundMessage(SenderMessage{
      openscreen::cast::SenderMessage::Type::kRpc,
      -1 /* sequence_number, unused by RPC messages */, true /* valid */,
      std::vector<uint8_t>(message.begin(), message.end())});
}

Error SenderSessionMessenger::SendInputMessage(ByteView message) {
  return SendOutboundMessage(SenderMessage{
      openscreen::cast::SenderMessage::Type::kInput,
      -1 /* sequence_number, unused by INPUT messages */, true /* valid */,
      std::vector<uint8_t>(message.begin(), message.end())});
}

Error SenderSessionMessenger::SendRequest(SenderMessage message,
                                          ReceiverMessage::Type reply_type,
                                          ReplyCallback cb) {
  // RPC and INPUT messages are not meant to be request/reply.
  OSP_CHECK(reply_type != ReceiverMessage::Type::kRpc);
  OSP_CHECK(reply_type != ReceiverMessage::Type::kInput);

  if (!cb) {
    return Error(Error::Code::kParameterInvalid,
                 "Must provide a reply callback");
  }
  const Error error = SendOutboundMessage(message);
  if (!error.ok()) {
    return error;
  }

  OSP_DCHECK(awaiting_replies_.find(message.sequence_number) ==
             awaiting_replies_.end());
  awaiting_replies_.emplace_back(message.sequence_number, std::move(cb));
  task_runner_->PostTaskWithDelay(
      [self = weak_factory_.GetWeakPtr(), seq_num = message.sequence_number] {
        if (self) {
          ReplyIfTimedOut(seq_num, &self->awaiting_replies_);
        }
      },
      kReplyTimeout);

  return Error::None();
}

void SenderSessionMessenger::OnMessage(const std::string& source_id,
                                       const std::string& message_namespace,
                                       const std::string& message) {
  if (source_id != receiver_id_ && source_id != kAnyDestination) {
    OSP_DLOG_WARN << "Received message from unknown/incorrect Cast Receiver "
                  << source_id << ". Currently connected to " << receiver_id_;
    return;
  }

  if (message_namespace != kCastWebrtcNamespace &&
      message_namespace != kCastRemotingNamespace) {
    OSP_DLOG_WARN << "Received message from unknown namespace: "
                  << message_namespace << ". Message was " << message;
    return;
  }

  ErrorOr<Json::Value> message_body = json::Parse(message);
  if (!message_body || !message_body.value().isObject()) {
    ReportError(message_body.error());
    OSP_DLOG_WARN << "Received an invalid message: " << message;
    return;
  }

  // If the message is valid JSON and we don't understand it, there are two
  // options: (1) it's an unknown type, or (2) the receiver filled out the
  // message incorrectly. In the first case we can drop it, it's likely just
  // unsupported. In the second case we might need it, so worth warning the
  // client.
  ErrorOr<ReceiverMessage> receiver_message =
      ReceiverMessage::Parse(message_body.value());
  if (receiver_message.is_error()) {
    ReportError(receiver_message.error());
    OSP_DLOG_WARN << "Received an invalid receiver message: "
                  << receiver_message.error();
    return;
  }

  if (receiver_message.value().type == ReceiverMessage::Type::kRpc) {
    if (rpc_callback_) {
      rpc_callback_(receiver_message.value());
    } else {
      OSP_DLOG_INFO << "Received RPC message but no callback, dropping";
    }
  } else if (receiver_message.value().type == ReceiverMessage::Type::kInput) {
    if (input_callback_) {
      input_callback_(receiver_message.value());
    } else {
      OSP_DLOG_INFO << "Received INPUT message but no callback, dropping";
    }
  } else {
    const int sequence_number = receiver_message.value().sequence_number;
    auto it = awaiting_replies_.find(sequence_number);
    if (it == awaiting_replies_.end()) {
      OSP_DLOG_WARN << "Received a reply I wasn't waiting for: "
                    << sequence_number;
      return;
    }

    ReplyCallback callback = std::move(it->second);
    awaiting_replies_.erase(it);
    callback(std::move(receiver_message.value()));
  }
}

void SenderSessionMessenger::OnError(const Error& error) {
  OSP_DLOG_WARN << "Received an error in the session messenger: " << error;
  ReportError(error);
}

ReceiverSessionMessenger::ReceiverSessionMessenger(MessagePort& message_port,
                                                   std::string source_id,
                                                   ErrorCallback cb)
    : SessionMessenger(message_port, std::move(source_id), std::move(cb)) {}

void ReceiverSessionMessenger::SetHandler(SenderMessage::Type type,
                                          RequestCallback cb) {
  OSP_DCHECK(callbacks_.find(type) == callbacks_.end());
  callbacks_.emplace_back(type, std::move(cb));
}

void ReceiverSessionMessenger::ResetHandler(SenderMessage::Type type) {
  callbacks_.erase_key(type);
}

Error ReceiverSessionMessenger::SendRpcMessage(const std::string& source_id,
                                               ByteView message) {
  return SendMessage(
      source_id,
      ReceiverMessage{ReceiverMessage::Type::kRpc, -1 /* sequence_number */,
                      true /* valid */,
                      std::vector<uint8_t>(message.begin(), message.end())});
}

Error ReceiverSessionMessenger::SendInputMessage(const std::string& source_id,
                                                 ByteView message) {
  return SendMessage(
      source_id,
      ReceiverMessage{ReceiverMessage::Type::kInput, -1 /* sequence_number */,
                      true /* valid */,
                      std::vector<uint8_t>(message.begin(), message.end())});
}

Error ReceiverSessionMessenger::SendMessage(const std::string& source_id,
                                            ReceiverMessage message) {
  if (source_id.empty()) {
    return Error(Error::Code::kInitializationFailure,
                 "Cannot send a message without a current source ID.");
  }

  const auto namespace_ = (message.type == ReceiverMessage::Type::kRpc ||
                           message.type == ReceiverMessage::Type::kInput)
                              ? kCastRemotingNamespace
                              : kCastWebrtcNamespace;

  ErrorOr<Json::Value> message_json = message.ToJson();
  OSP_CHECK(message_json.is_value()) << "Tried to send an invalid message";
  return SessionMessenger::SendMessage(source_id, namespace_,
                                       message_json.value());
}

void ReceiverSessionMessenger::SetCustomMessageHandler(
    std::string_view message_namespace,
    CustomMessageCallback cb) {
  auto it = std::find_if(custom_message_handlers_.begin(),
                         custom_message_handlers_.end(),
                         [&message_namespace](const auto& pair) {
                           return pair.first == message_namespace;
                         });

  if (!cb) {
    if (it != custom_message_handlers_.end()) {
      custom_message_handlers_.erase(it);
    }
    return;
  }

  if (it != custom_message_handlers_.end()) {
    OSP_LOG_ERROR << "Handler already exists for namespace: "
                  << message_namespace;
    return;
  } else {
    custom_message_handlers_.emplace_back(std::string(message_namespace),
                                          std::move(cb));
  }
}

Error ReceiverSessionMessenger::SendMessage(std::string_view destination_id,
                                            std::string_view message_namespace,
                                            std::string_view message) {
  message_port().PostMessage(std::string(destination_id),
                             std::string(message_namespace),
                             std::string(message));
  return Error::None();
}

void ReceiverSessionMessenger::OnMessage(const std::string& source_id,
                                         const std::string& message_namespace,
                                         const std::string& message) {
  if (message_namespace != kCastWebrtcNamespace &&
      message_namespace != kCastRemotingNamespace) {
    auto it = std::find_if(custom_message_handlers_.begin(),
                           custom_message_handlers_.end(),
                           [&message_namespace](const auto& pair) {
                             return pair.first == message_namespace;
                           });
    if (it != custom_message_handlers_.end()) {
      it->second(source_id, message_namespace, message);
      return;
    }
    OSP_DLOG_WARN << "Received message from unknown namespace: "
                  << message_namespace;
    return;
  }

  // If the message is bad JSON, the sender is in a funky state so we
  // report an error.
  ErrorOr<Json::Value> message_body = json::Parse(message);
  if (message_body.is_error() || !message_body.value().isObject()) {
    ReportError(message_body.error());
    return;
  }

  // If the message is valid JSON and we don't understand it, there are two
  // options: (1) it's an unknown type, or (2) the sender filled out the message
  // incorrectly. In the first case we can drop it, it's likely just
  // unsupported. In the second case we might need it, so worth warning the
  // client.
  ErrorOr<SenderMessage> sender_message =
      SenderMessage::Parse(message_body.value());
  if (sender_message.is_error()) {
    ReportError(sender_message.error());
    OSP_DLOG_WARN << "Received an invalid sender message: "
                  << sender_message.error();
    return;
  }

  if (sender_message.value().type == SenderMessage::Type::kOffer ||
      sender_message.value().type == SenderMessage::Type::kGetCapabilities) {
    OSP_VLOG << "Received Message:\n" << message;
  }

  auto it = callbacks_.find(sender_message.value().type);
  if (it == callbacks_.end()) {
    OSP_DLOG_INFO << "Received message without a callback, dropping";
    return;
  }
  it->second(source_id, sender_message.value());
}

void ReceiverSessionMessenger::OnError(const Error& error) {
  OSP_DLOG_WARN << "Received an error in the session messenger: " << error;
  ReportError(error);
}

}  // namespace openscreen::cast