#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"

#define MAXCHUNKSIZE (65536+16)

// 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 enclace-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 MAXCHUNKSIZE 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;
    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[MAXCHUNKSIZE];

    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);

    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,
        std::function<void(uint32_t)> message_cb,
        std::function<void(uint8_t*,uint32_t)> chunk_cb);
};

class NetIO {
    const Config &conf;
    const NodeConfig &myconf;
    std::deque<std::optional<NodeIO>> nodeios;

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

    size_t num_nodes;
    size_t me;
    NodeIO &node(size_t node_num) {
        assert(node_num < num_nodes);
        return nodeios[node_num].value();
    }
};

extern NetIO *g_netio;

#endif