#ifndef NETWORK_BUFFER_HH
#define NETWORK_BUFFER_HH

/**
 * @file
 *
 * Buffering of network streams
 */

#include "Socket.hh"
#include "../Error.hh"

#include <cassert>

/**
 * Minimum chunk size to avoid handling single bytes at a time (for resize, mainly)
 */
const size_t NETWORK_BUFFER_CHUNK_SIZE = 1024;

/**
 * Base class of errors thrown by NetworkBuffer* methods
 */
class NetworkBufferError : public Error {
    public:
        NetworkBufferError (const std::string &message) : Error(message) { }
};

/**
 * Base buffer-manipulation operations for buffered socket send/recv.
 *
 * This implements a simple linear in-memory buffer which can contain data that must be buffered for send/recv.
 * New data can be written to the end of the buffer, and data can be read/discarded from the beginning.
 *
 * A buffer is associated with a NetworkSocket, and takes care of all send/recv activity.
 */
class NetworkBufferBase {
    protected:
        /** The socket that we use */
        NetworkSocket *socket;
        
        /** The buffer itself */
        char *buf;

        /** Buffer size and current write offset */
        size_t size, offset;
    
    public:
        /**
         * Allocate buf using the given initial size, and set offset to zero
         */
        NetworkBufferBase (NetworkSocket *socket, size_t size_hint);

        /**
         * Free()'s the buf
         */
        ~NetworkBufferBase (void);
    
    private:
        /**
         * No copying, these are undefined
         */
        NetworkBufferBase (const NetworkBufferBase &copy);
        NetworkBufferBase& operator= (const NetworkBufferBase &copy);
    
    protected:
        /**
         * Resize the buffer, allocating enough new space to hold <item_size> bytes at the end, and leaving any
         * existing data at the beginning in-place.
         *
         * This is done by doubling the buffer size until item_size fits, and then reallocing (if needed).
         * Currently, the buffer is never shrunk, although it probably should be to adapt better to peak load (e.g.
         * when a player joins and the map data is transmitted...).
         *
         * @param item_size the number of bytes that must fit at the end of the buffer
         */
        void resize (size_t item_size);

        /**
         * Trim the buffer, discarding <prefix_size> bytes at the beginning. Updates offset to match.
         *
         * Currently, this is implemented as a simple memmove() of all the remaining data in the buffer, which may be
         * slow under some load patterns.
         *
         * @param prefix_size the number of bytes to discard
         */
        void trim (size_t prefix_size);
};

/**
 * Buffered prefix-len socket input
 *
 * This handles calling recv() on the socket, and is specialized to handle a stream of messages prefixed with a length
 * header (uint16_t or uint32_t). Use the peek_data() method to receive these messages when NetworkSocket::sig_read() 
 * indicates that more data is available.
 */
class NetworkBufferInput : public NetworkBufferBase {
    public:
        /**
         * @see NetworkBufferBase
         */
        NetworkBufferInput (NetworkSocket *socket, size_t size_hint);

    private:
        /**
         * Attempts to recv the given number of bytes into our buf, returning true on success
         *
         * @param item_size minimum number of bytes of data that we need in the buffer
         * @return bool true if the buffer now contains at least item_size bytes
         */
        bool try_read (size_t item_size);

        /**
         * Tests if the buffer contains at least the given amount of data, but doesn't recv or anything
         *
         * @param data_size number of bytes that we are expecting
         * @return bool true if the buffer contains at least data_size bytes
         */
        bool have_data (size_t data_size);
    
    public:
        // @{
        /**
         *
         * Attempts to read the length prefix into val_ref, returning true on success, false if there's not enough data
         * in the buffer
         *
         * @param val_ref stores the value read here if we have it
         * @return bool was val_ref set
         */
        bool peek_prefix (uint16_t &val_ref);
        bool peek_prefix (uint32_t &val_ref);
        // @} 
        
        /**
         * This attempts to locate a full message in the buffer (calling recv as nessecary), prefixed with a length
         * prefix of the given type (using peek_prefix). If succesfull (we have a full message in the buffer), 
         * \a prefix will be updated to the length of the message (not including the prefix), and \a buf_ref will be
         * updated to point at the message of \a prefix bytes (not including the prefix) in our internal buffer memory.
         * If a full message could not be received (recv would block), this will return false.
         *
         * Once you have processed the message, call flush_data() to remove the unused message from the buffer.
         *
         * @param prefix updated to the message length
         * @param buf_ref updated to point at the message data
         * @return bool true if we have a full message, false if we need to wait for more data on the socket
         *
         * @see peek_prefix
         * @see flush_data
         */
        template <typename PrefixType> bool peek_data (PrefixType &prefix, char *&buf_ref);
          
        /**
         * This flushes the current message from the buffer, as returned when peek_data returns true. It is a bug to
         * call flush_data when there is no full message present, the behaviour is unspecified (most likely an
         * assert()).
         *
         * @see peek_data
         */
        template <typename PrefixType> void flush_data (void);
};

/**
 * Buffered prefix-len socket output.
 *
 * This handles calling send() on the socket, and is specialized to handle a stream of message prefixed with a length
 * header (uint16_t or uint32_t). You can write messages using write_prefix(), and they will be buffered if needed.
 * If write_prefix() returns true (socket buffer full), then you must register the socket for
 * NetworkSocket::set_poll_write(), and call flush_write() once NetworkSocket::sig_write() is triggered. The socket's
 * poll_write should be unregistered once flush_write returns false (no more buffered data remaining to be sent), or
 * an many processor cycles will be wasted, for the socket will remain ready for write until its buffer fills up again.
 *
 * Data is not buffered needlessly; if the socket's buffer has room, write_prefix will not have to touch our internal
 * buffer. In fact, write_prefix will rarely return false except under heavy network congestion with high levels of
 * traffic on the socket.
 */
class NetworkBufferOutput : public NetworkBufferBase {
    public:
        /**
         * @see NetworkBufferBase
         */
        NetworkBufferOutput (NetworkSocket *socket, size_t size_hint);

    private:
        /**
         * Write the given data to the socket, either now or later. 
         *
         * If our buffer is empty, fast-path the given buf_ptr directly to send(), then copy the remaining portion to
         * our buffer for later use with flush_write.
         *
         * @param buf_ptr the data that we need to send
         * @param buf_size number of bytes in buf_ptr
         */
        void push_write (char *buf_ptr, size_t buf_size);
   
    public:    
        /**
         * If we have data in our buffer, flush it out using send(). This should be called once the socket indicates it
         * is ready for write again after a call to write_prefix that returned false. Once this returns false, the
         * NetworkSocket::set_poll_write() flag should be unset again to avoid needless calls to this.
         *
         * @return true if there's still buffered data left to write, false otherwise (buffer is empty)
         */
        bool flush_write (void);
        
        // @{
        /**
         * Write out the given message, writing first the prefix, and then the data itself, using push_write.
         *
         * Returns true if the data was passed on to the socket API directly, false if a portion of it had to be
         * buffered (if there is already data buffered, all subsequent write_prefix's will buffer the full message
         * until our buffer is empty again). If this method returns false, you must register the socket for
         * poll_write and call flush_write later.
         * 
         * @param buf the data to write
         * @param prefix the amount of data
         * @return true if we had to buffer data, false otherwise
         */
        bool write_prefix (char *buf, uint16_t prefix);
        bool write_prefix (char *buf, uint32_t prefix);
        // @}
};


/*
 * NetworkBufferInput template method implementation
 */
template <typename PrefixType> bool NetworkBufferOutput::peek_data (PrefixType &prefix, char *&buf_ref) {
    size_t missing = 0;
    
    do {    
        // do we have the prefix?
        if (peek_prefix(prefix)) {
            // do we already have the payload?
            if (offset >= sizeof(PrefixType) + prefix) {
                break;

            } else {
                missing = (sizeof(PrefixType) + prefix) - offset;
            }

        } else {
            missing = sizeof(PrefixType);
        }

        // sanity-check (above >= and - should never happen like this)
        assert(missing);
        
        // try and read the missing data
        if (try_read(missing) == false) {
            // if unable to read what we need, return zero.
            return false;
        }
        
        // assess the situation again
    } while (true);
    
    // update the buf_ref to point past the prefix-length
    buf_ref = buf + sizeof(PrefixType);

    // return message
    return true;
}

template <typename PrefixType> void NetworkBufferOutput::flush_data (void) {
    PrefixType prefix;
    
    // we *must* have a valid prefix
    if (!peek_prefix(prefix))
        assert(false);

    // ensure that we have the data to trim...
    assert(offset >= sizeof(PrefixType) + prefix);

    // trim the bytes out
    trim(sizeof(PrefixType) + prefix);
}

#endif