#ifndef BSTREAM_H
#define BSTREAM_H

#include <stdlib.h>
#include <string.h>
#include <stdexcept>
#include <winsock2.h>
#include <string>

class BOutputStream
{
public:
	BOutputStream()
	{
		m_ptr = (char*) malloc( m_capacity = 1500 );
		CheckPtr();
		m_size = 0;
	}
	~BOutputStream()
	{
		free( m_ptr	);
	}

	void Rewind( size_t sz )
	{
		if (sz > m_size)
			throw std::runtime_error("Attempt to rewind past start of stream");
		m_size -= sz;
	}

	void * Put( size_t sz )
	{
		CheckPtr();
		while (sz + m_size > m_capacity)
			IncCap();
		void * out = m_ptr + m_size;
		m_size += sz;
		return out;
	}

	void Put( const void * ptr, size_t sz )
	{
		memcpy( Put(sz), ptr, sz );
	}

	const void * GetPtr() const { return m_ptr; }
	const size_t GetSize() const { return m_size; }

	void Clear() { m_size = 0; }

private:
	BOutputStream( const BOutputStream& );
	BOutputStream& operator=( const BOutputStream& );

	char * m_ptr;
	size_t m_capacity;
	size_t m_size;

	void CheckPtr()
	{
		if (!m_ptr)
			throw std::runtime_error( "Allocation failure" );
	}

	void IncCap()
	{
		m_ptr = (char*) realloc( m_ptr, m_capacity *= 2 );
		CheckPtr();
	}
};

inline BOutputStream& operator<<( BOutputStream& out, char c )
{
	*(char*)out.Put(sizeof(char)) = c;
	return out;
}

inline BOutputStream& operator<<( BOutputStream& out, unsigned char c )
{
	*(unsigned char*)out.Put(sizeof(unsigned char)) = c;
	return out;
}

inline BOutputStream& operator<<( BOutputStream& out, unsigned short x )
{
	*(unsigned short*)out.Put(sizeof(unsigned short)) = htons(x);
	return out;
}

inline BOutputStream& operator<<( BOutputStream& out, unsigned x )
{
	*(unsigned*)out.Put(sizeof(unsigned)) = htonl(x);
	return out;
}

inline BOutputStream& operator<<( BOutputStream& out, const std::string& str )
{
	if (str.length() > 255)
		throw std::runtime_error("String too long: " + str);
	out << (unsigned char)str.length();
	for (size_t i=0; i<str.length(); i++)
		out << str[i];
	return out;
}

class BInputStream
{
public:
	BInputStream( const void * start, size_t length )
	{
		m_input = (const char *) start;
		m_end = m_input + length;
	}

	const void * Get( size_t sz )
	{
		const char * out = m_input;
		m_input += sz;
		if (m_input > m_end)
			throw std::runtime_error("Attempt to read too far");
		return out;
	}

	void Get( void * buf, size_t sz )
	{
		memcpy( buf, Get(sz), sz );
	}

	size_t Remaining() const
	{
		return m_end - m_input;
	}

private:
	const char * m_input;
	const char * m_end;
};

inline BInputStream& operator>>( BInputStream& in, char& x )
{
	x = *(char*)in.Get(sizeof(char));
	return in;
}

inline BInputStream& operator>>( BInputStream& in, unsigned char& x )
{
	x = *(unsigned char*)in.Get(sizeof(unsigned char));
	return in;
}

inline BInputStream& operator>>( BInputStream& in, unsigned short& x )
{
	x = ntohs(*(unsigned short*)in.Get(sizeof(unsigned short)));
	return in;
}

inline BInputStream& operator>>( BInputStream& in, unsigned& x )
{
	x = ntohl(*(unsigned*)in.Get(sizeof(unsigned)));
	return in;
}

inline BInputStream& operator>>( BInputStream& in, std::string& str )
{
	unsigned char l;
	in >> l;
	str.resize(0);
	while (l--)
	{
		char c;
		in >> c;
		str += c;
	}
	return in;
}

#endif
