// this is for emacs file handling -*- mode: c++; indent-tabs-mode: nil -*-

// -- BEGIN LICENSE BLOCK ----------------------------------------------
// Copyright 2021 FZI Forschungszentrum Informatik
// Created on behalf of Universal Robots A/S
//
// 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.
// -- END LICENSE BLOCK ------------------------------------------------

//----------------------------------------------------------------------
/*!\file
 *
 * \author  Felix Exner mauch@fzi.de
 * \date    2021-03-13
 *
 */
//----------------------------------------------------------------------

#ifndef UR_CLIENT_LIBRARY_TCP_SERVER_H_INCLUDED
#define UR_CLIENT_LIBRARY_TCP_SERVER_H_INCLUDED

#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <mutex>
#include <thread>
#include <vector>

#include "ur_client_library/comm/socket_t.h"

namespace urcl
{
namespace comm
{
/*!
 * \brief Wrapper class for a TCP socket server.
 *
 * The server can be created given a port and it can register callbacks for 3 events:
 *  - connect: A new client connected
 *  - disconnect: A client disconnected
 *  - message: Data sent from one of the clients
 *
 *  Please note that start() has to be called manually after initialization and callback
 *  registering in order to start handling socket events.
 *
 *  While this server implementation supports multiple (number limited by system's socket
 *  implementation) clients by default, a maximum number of allowed clients can be configured.
 */
class TCPServer
{
public:
  TCPServer() = delete;
  /*!
   * \brief Create a TCPServer object
   *
   * \param port Port on which to operate. The port will be bound to the process creating the
   * object.
   * \param max_num_tries If binding the socket fails, it will be retried this many times. If 0 is
   * specified, binding the socket will be tried indefinitely.
   * \param reconnection_time Wait time in between binding attempts.
   */
  explicit TCPServer(const int port, const size_t max_num_tries = 0,
                     const std::chrono::milliseconds reconnection_time = std::chrono::seconds(1));
  virtual ~TCPServer();

  /*!
   * \brief This callback will be triggered on clients connecting to the server
   *
   * \param func Function handling the event information. The file descriptor created by the
   * connection event will be passed to the function.
   *
   * \note: The connection callback will be triggered with the socket being accepted. Hence, it
   * is possible to send data from the connection callback directly.
   */
  void setConnectCallback(std::function<void(const socket_t)> func)
  {
    std::lock_guard<std::mutex> lk(callback_mutex_);
    new_connection_callback_ = func;
  }

  /*!
   * \brief This callback will be triggered on clients disconnecting from the server
   *
   * \param func Function handling the event information. The file descriptor created by the
   * connection event will be passed to the function.
   *
   * \note: The socket will already be closed when the disconnect callback is triggered, thus
   * trying to interact with the socket from the disconnect callback will fail.
   */
  void setDisconnectCallback(std::function<void(const socket_t)> func)
  {
    std::lock_guard<std::mutex> lk(callback_mutex_);
    disconnect_callback_ = func;
  }

  /*!
   * \brief This callback will be triggered on messages received on the socket
   *
   * \param func Function handling the event information. The file client's file_descriptor will be
   * passed to the function as well as the actual message received from the client.
   */
  void setMessageCallback(std::function<void(const socket_t, char*, int)> func)
  {
    std::lock_guard<std::mutex> lk(message_mutex_);
    message_callback_ = func;
  }

  /*!
   * \brief Start event handling.
   *
   * Without calling this function the socket will be advertised and bound to a tcp port, but no
   * handling of connection requests will be performed.
   */
  void start();

  /*!
   * \brief Shutdown the server and close all client connections.
   *
   * \note: This should not be called from within any of the registered callback functions, as
   * it will cause a deadlock. If you want to shutdown the server from a callback, you can e.g.
   * start a new thread that calls shutdown() from there.
   */
  void shutdown();

  /*!
   * \brief Writes to a client
   *
   * \param[in] fd File descriptor belonging to the client the data should be sent to. The file
   * descriptor will be given from the connection callback.
   * \param[in] buf Buffer of bytes to write
   * \param[in] buf_len Number of bytes in the buffer
   * \param[out] written Number of bytes actually written
   *
   * \returns True on success, false otherwise
   */
  bool write(const socket_t fd, const uint8_t* buf, const size_t buf_len, size_t& written);

  /*!
   * \brief Writes to a filedescriptor without verifying that it is a client or even a valid
   * filedescriptor. It is the caller's responsibility to ensure that the filedescriptor is valid
   * and belongs to a client.
   *
   * \param[in] fd File descriptor belonging to the client the data should be sent to. The file
   * descriptor will be given from the connection callback.
   * \param[in] buf Buffer of bytes to write
   * \param[in] buf_len Number of bytes in the buffer
   * \param[out] written Number of bytes actually written
   *
   * \returns True on success, false otherwise
   */
  bool writeUnchecked(const socket_t fd, const uint8_t* buf, const size_t buf_len, size_t& written);

  /*!
   * \brief Get the maximum number of clients allowed to connect to this server
   *
   * \returns The currently configured client limit. 0 means unlimited amount of clients allowed.
   */
  uint32_t getMaxClientsAllowed() const
  {
    return max_clients_allowed_;
  }

  /*!
   * \brief Set the maximum number of clients allowed to connect to this server.
   *
   * 0 means unlimited number of clients allowed.
   *
   */
  void setMaxClientsAllowed(const uint32_t& max_clients_allowed)
  {
    max_clients_allowed_ = max_clients_allowed;
  }

  /*!
   * \brief Get the port this server is bound to
   *
   * If port number 0 is passed during initialization, the server will bind to a random free port.
   * In this case, this function can be used to get the actual port number the server is bound to.
   * This should only be called after the server has been initialized, otherwise the returned port
   * number might not be correct.
   *
   * \returns The port number this server is bound to
   */
  int getPort() const
  {
    return port_;
  }

private:
  void init();
  void bind(const size_t max_num_tries, const std::chrono::milliseconds reconnection_time);
  void startListen();

  //! Handles connection events
  void handleConnect();

  void handleDisconnect(const socket_t fd);

  //! read data from socket
  bool readData(const socket_t fd);

  //! Event handler. Blocks until activity on any client or connection attempt
  void spin();

  //! Runs spin() as long as keep_running_ is set to true.
  void worker();

  std::atomic<bool> keep_running_{ false };
  std::thread worker_thread_;

  std::atomic<socket_t> listen_fd_;
  int port_;

  socket_t maxfd_;
  fd_set masterfds_;
  fd_set tempfds_;

  uint32_t max_clients_allowed_;
  std::vector<socket_t> client_fds_;
  std::mutex clients_mutex_;
  std::mutex message_mutex_;
  std::mutex listen_fd_mutex_;
  std::mutex callback_mutex_;

  static const int INPUT_BUFFER_SIZE = 4096;
  char input_buffer_[INPUT_BUFFER_SIZE];

  std::function<void(const socket_t)> new_connection_callback_;
  std::function<void(const socket_t)> disconnect_callback_;
  std::function<void(const socket_t, char* buffer, int nbytesrecv)> message_callback_;
};

}  // namespace comm
}  // namespace urcl

#endif  // ifndef UR_CLIENT_LIBRARY_TCP_SERVER_H_INCLUDED
