#ifndef __NET_HPP__
#define __NET_HPP__

#include <vector>
#include <deque>
#include <optional>
#include <functional>
#include <tuple>
#include <boost/asio.hpp>
#include <boost/thread.hpp>

#include "appconfig.hpp"
#include "../Enclave/enclave_api.h"

// The inter-node (untrusted node to untrusted node) communication
// protocol is as follows.  Nodes are numbered 0 through num_nodes-1.
// At startup time, each pair of nodes establishes a TCP connection by
// having the lower-numbered node connect to the higher-numbered node,
// and send a two-byte value of its (the sender's) node number.  Once
// all the connections are established, commands consist of a 5-byte
// header, followed optionally by some data.  The commands are listed
// below.  If a socket closes, we interpret that to mean the experiment
// is over, and the node shuts down (which will close its own sockets,
// its peers will shut down, etc.).  [This isn't the best idea for a
// robust long-lived deployment, of course.]
//
// The commands are:
//
// EPOCH: 0x00 + 4-byte epoch number (little-endian)
//
// This command is sent by the leader (typically node 0) to each other
// node at the start of each epoch.
//
// MESSAGE: 0x01 + 4-byte total message length (little-endian)
//
// This command says that a number of CHUNKs comprising a single
// enclave-to-enclave message will follow, whose total size will be the
// given value.  Note that the data itself is sent following a CHUNK
// header, not a MESSAGE header, even if it's small.
//
// CHUNK: 0x02 + 4-byte chunk length (little-endian)
// + that many bytes of data
//
// This command transmits the enclave-to-enclave data.  The data in the
// chunk will be (after the enclave-to-enclave handshake, anyway)
// AES-GCM encrypted to a key known to the receiving enclave (but not
// the receiving untrusted node).  The chunk number (starting from 0 and
// not reset between messages) will be the IV, which is not transmitted.
// The 16-byte GCM tag will be the last 16 bytes of the chunk (and
// included in the length in the chunk header).  The sum of the chunk
// lengths since the last MESSAGE command may not exceed the length in
// that MESSAGE command.

// Data for chunks are stored in frames.  The frames are pre-allocated
// to be FRAME_SIZE bytes each, and reused as much as possible by the
// NodeIO class.  A node will request a frame from the NodeIO, which
// will return a pointer.  The node will pass that pointer to the
// enclave, which will write data into it, and also return to the node
// how much data it wrote.  The node will async_write the chunk header
// and the chunk data.  The async write completion handler will return
// the frame to the NodeIO when the write completes.
//
// Headers are stored as the low 5 bytes of a uint64_t.  Note that means
// for headers containing sizes, the value of this uint64_t will be (for
// example for the CHUNK header) (chunk_len << 8) + 0x02.

using boost::asio::ip::tcp;

class NodeIO {
    tcp::socket sock;
    nodenum_t node_num;
    using CommandTuple = std::tuple<uint64_t,uint8_t*,size_t>;
    std::deque<CommandTuple> commands_inflight;
    std::deque<uint8_t *> frames_available;
    // The frames and commands are used and returned by different
    // threads, so we protect them with a mutex each
    boost::mutex frame_deque_lock, commands_deque_lock;

    // The claimed size of the message currently being sent in chunks
    uint32_t msgsize_inflight;
    // The total size of the chunks so far we've sent for this message
    uint32_t chunksize_inflight;

    // As above, but for incoming messages and chunks
    uint32_t recv_msgsize_inflight;
    uint32_t recv_chunksize_inflight;

    // The static uint64_t used to receive a header
    uint64_t receive_header;
    // The static frame used to receive a chunk
    uint8_t receive_frame[FRAME_SIZE];

    uint64_t bytes_sent; // count bytes sent

    void send_header_data(uint64_t header, uint8_t *data, size_t len);

    // Asynchronously send the first message from the command queue.
    // * The command_deque_lock must be held when this is called! *
    // This method may be called from either thread (the work thread or
    // the async_write handler thread).
    void async_send_commands();

public:
    NodeIO(tcp::socket &&socket, nodenum_t node_num);
    uint8_t *request_frame();
    void return_frame(uint8_t* frame);

    void send_epoch(uint32_t epoch_num);
    void send_message_header(uint32_t tot_message_len);
    // Returns true if there are more chunks to send in this message,
    // false if not.
    bool send_chunk(uint8_t *data, uint32_t chunk_len);

    // Asynchronously receive commands from this socket.  Depending on
    // what they are, one of the three callbacks will be called.  The
    // callbacks may be called from a different thread.  The data
    // pointer in chunk_cb is to a _static_ frame that's only used for
    // receiving.  Be sure to do whatever you need to do with the
    // contents (typically, pass it to the enclave) before calling this
    // function again.
    void recv_commands(
        std::function<void(boost::system::error_code)> error_cb,
        std::function<void(uint32_t)> epoch_cb);

    // Close the socket
    void close() { sock.close(); }

    uint64_t reset_bytes_sent();
};

class NetIO {
    boost::asio::io_context &context;
    const Config &conf;
    const NodeConfig &myconf;
    std::deque<std::optional<NodeIO>> nodeios;
    std::shared_ptr<tcp::acceptor> ingestion_acceptor;
    std::shared_ptr<tcp::acceptor> storage_acceptor;

    size_t auth_size, msgbundle_size;
    void ing_receive_msgbundle(tcp::socket* socket, clientid_t c_simid);
    void ing_authenticate_new_client(tcp::socket* socket,
        const boost::system::error_code& error);
    void ing_start_accept();

    std::vector<tcp::socket*> client_sockets;
    uint32_t num_clients_per_stg;
    unsigned char *epoch_tokens;
    unsigned char *epoch_mailboxes;
    uint16_t num_stg_nodes;
    uint32_t token_bundle_size;
    uint32_t mailbox_size;
    void stg_authenticate_new_client(tcp::socket* socket,
        const boost::system::error_code& error);
    void stg_start_accept();

public:
    NetIO(boost::asio::io_context &io_context, const Config &config);

    nodenum_t num_nodes;
    nodenum_t me;
    NodeIO &node(nodenum_t node_num) {
        assert(node_num < num_nodes);
        return nodeios[node_num].value();
    }
    const Config &config() { return conf; }
    const NodeConfig &myconfig() { return myconf; }
    boost::asio::io_context &io_context() { return context; }
    // Call recv_commands with these arguments on each of the nodes (not
    // including ourselves)
    void recv_commands(
        std::function<void(boost::system::error_code)> error_cb,
        std::function<void(uint32_t)> epoch_cb);

    void send_client_mailbox();

    // Close all the sockets
    void close();

    uint64_t reset_bytes_sent();

#ifdef TRACE_SOCKIO
private:
    struct timeval last_ing;
    size_t num_ing;
#endif
};

extern NetIO *g_netio;
extern size_t client_count;

#endif