#ifndef NETWORK_BUFFER_HH
#define NETWORK_BUFFER_HH

#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
 */
class NetworkBufferBase {
    protected:
        /** The socket that we use */
        NetworkSocket socket;
        
        /** The buffer itself */
        char *buf;

        /** Buffer size and current read/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
         *
         * @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.
         *
         * @param prefix_size the number of bytes to discard
         */
        void trim (size_t prefix_size);
};

/**
 * Buffered prefix-len socket input
 */
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 read a length-prefix of the given type (using peek_prefix), and then the associated data.
         * If succesful, this sets prefix to the length of the data, and buf_ref to point at the data inside our buffer
         * and returns true, else false.
         *
         * This will try and consume data from the buffer, or recv if needed.
         *
         * @param prefix stores the data length here
         * @param buf_ref stores a pointer to the data here
         * @return bool true if we have the full data, 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) {
            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
                // XXX: a zero-prefix will trigger 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
            return true;
        }
          
        /**
         * This flushes a prefix-length worth of data from the buffer, i.e. it first reads the prefix, and then trims
         * the prefix and the data away. Don't call this unless you *know* that the buffer contains enough data.
         *
         * @see peek_data
         */
        template <typename PrefixType> void flush_data (void) {
            PrefixType prefix;
            
            // we *must* have a valid prefix
            if (!peek_prefix(prefix))
                assert(false);

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

/**
 * Buffered prefix-len socket output
 */
class NetworkBufferOutput : public NetworkBufferBase {
    public:
        /**
         * @see NetworkBufferBase
         */
        NetworkBufferOutput (NetworkSocket &socket, size_t size_hint);

    private:
        /**
         * Write the given data to the socket, either now of 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().
         */
        void flush_write (void);
        
        // @{
        /**
         * Write out the given data, writing first the prefix, and then the data itself, using push_write.
         *
         * @param buf the data to write
         * @param prefix the amount of data
         */
        void write_prefix (char *buf, uint16_t prefix);
        void write_prefix (char *buf, uint32_t prefix);
        // @}
};

#endif