#!/usr/bin/env python

import socket
import gzip
try:
	from cStringIO import StringIO
except ImportError, e:
	from StringIO import StringIO
from struct import unpack, pack
from __init__ import dumps, loads

def _bintoint(data):
	return unpack("<i", data)[0]

def _sendobj(self, obj):
	"""
	Atomically send a BSON message.
	"""
	data = dumps(obj)
	self.sendall(data)


def _recvobj(self):
	"""
	Atomic read of a BSON message.

	This function either returns a dict, None, or raises a socket error.

	If the return value is None, it means the socket is closed by the other side.
	"""
	sock_buf = self.recvbytes(4)
	if sock_buf is None:
		return None

	message_length = _bintoint(sock_buf.getvalue())
	sock_buf = self.recvbytes(message_length - 4, sock_buf)
	if sock_buf is None:
		return None

	retval = loads(sock_buf.getvalue())
	return retval


def _gzip_sendobj(self, obj):
	data = dumps(obj)
	buf = StringIO()
	fd = gzip.GzipFile(mode="wb", fileobj=buf)
	fd.write(data)
	fd.close()
	zip_data = buf.getvalue()
	self.sendall(pack('<i', len(zip_data)) + zip_data)


def _gzip_recvobj(self):
	sock_buf = self.recvbytes(4)
	if sock_buf is None:
		return None
	message_length = _bintoint(sock_buf.getvalue())
	sock_buf = self.recvbytes(message_length)
	if sock_buf is None:
		return None
	
	sock_buf.seek(0)
	fd = gzip.GzipFile(mode="rb", fileobj=sock_buf)
	data = fd.read()
	retval = loads(data)
	return retval


def _recvbytes(self, bytes_needed, sock_buf = None):
	"""
	Atomic read of bytes_needed bytes.

	This function either returns exactly the nmber of bytes requested in a
	StringIO buffer, None, or raises a socket error.

	If the return value is None, it means the socket is closed by the other side.
	"""
	if sock_buf is None:
		sock_buf = StringIO()
	bytes_count = 0
	while bytes_count < bytes_needed:
		chunk = self.recv(min(bytes_needed - bytes_count, 32768))
		part_count = len(chunk)

		if part_count < 1:
			return None

		bytes_count += part_count
		sock_buf.write(chunk)
	
	return sock_buf
