#!/usr/bin/python -OOOO
# vim: set fileencoding=utf8 shiftwidth=4 tabstop=4 textwidth=80 foldmethod=marker :
# Copyright (c) 2010, Kou Man Tong. All rights reserved.
# For licensing, see LICENSE file included in the package.
"""
Base codec functions for bson.
"""
import struct
import cStringIO
import calendar, pytz
from datetime import datetime
import warnings
try:
	from abc import ABCMeta, abstractmethod
except ImportError:
	# If abc is not present (older versions of python), just define the ABCMeta
	# class as a dummy class, we don't really need it anyway.
	class ABCMeta(type):
		pass

	def abstractmethod(funcobj):
		"""A decorator indicating abstract methods.

		Requires that the metaclass is ABCMeta or derived from it.	A
		class that has a metaclass derived from ABCMeta cannot be
		instantiated unless all of its abstract methods are overridden.
		The abstract methods can be called using any of the normal
		'super' call mechanisms.

		Usage:

			class C:

			__metaclass__ = ABCMeta

			@abstractmethod
			def my_abstract_method(self, ...):
				...
		"""
		funcobj.__isabstractmethod__ = True
		return funcobj


# {{{ Error Classes
class MissingClassDefinition(ValueError):
	def __init__(self, class_name):
		super(MissingClassDefinition, self).__init__(
		"No class definition for class %s" % (class_name,))
# }}}
# {{{ Warning Classes
class MissingTimezoneWarning(RuntimeWarning):
	def __init__(self, *args):
		

# }}}
# {{{ Traversal Step
class TraversalStep(object):
	def __init__(self, parent, key):
		self.parent = parent
		self.key = key
# }}}
# {{{ Custom Object Codec

class BSONCoding(object):
	__metaclass__ = ABCMeta

	@abstractmethod
	def bson_encode(self):
		pass

	@abstractmethod
	def bson_init(self, raw_values):
		pass

classes = {}

def import_class(cls):
	if not issubclass(cls, BSONCoding):
		return

	global classes
	classes[cls.__name__] = cls

def import_classes(*args):
	for cls in args:
		import_class(cls)

def import_classes_from_modules(*args):
	for module in args:
		for item in module.__dict__:
			if hasattr(item, "__new__") and hasattr(item, "__name__"):
				import_class(item)

def encode_object(obj, traversal_stack, generator_func):
	values = obj.bson_encode()
	class_name = obj.__class__.__name__
	values["$$__CLASS_NAME__$$"] = class_name
	return encode_document(values, traversal_stack, obj, generator_func)

def encode_object_element(name, value, traversal_stack, generator_func):
	return "\x03" + encode_cstring(name) + \
			encode_object(value, traversal_stack,
					generator_func = generator_func)

class _EmptyClass(object):
	pass

def decode_object(raw_values):
	global classes
	class_name = raw_values["$$__CLASS_NAME__$$"]
	cls = None
	try:
		cls = classes[class_name]
	except KeyError, e:
		raise MissingClassDefinition(class_name)

	retval = _EmptyClass()
	retval.__class__ = cls
	alt_retval = retval.bson_init(raw_values)
	return alt_retval or retval

# }}}
# {{{ Codec Logic
def encode_string(value):
	value = value.encode("utf8")
	length = len(value)
	return struct.pack("<i%dsb" % (length,), length + 1, value, 0)

def decode_string(data, base):
	length = struct.unpack("<i", data[base:base + 4])[0]
	value = data[base + 4: base + 4 + length - 1]
	value = value.decode("utf8")
	return (base + 4 + length, value)

def encode_cstring(value):
	if isinstance(value, unicode):
		value = value.encode("utf8")
	return value + "\x00"

def decode_cstring(data, base):
	length = 0
	max_length = len(data) - base
	while length < max_length:
		character = data[base + length]
		length += 1
		if character == "\x00":
			break
	return (base + length, data[base:base + length - 1].decode("utf8"))

def encode_binary(value):
	length = len(value)
	return struct.pack("<ib", length, 0) + value

def decode_binary(data, base):
	length, binary_type = struct.unpack("<ib", data[base:base + 5])
	return (base + 5 + length, data[base + 5:base + 5 + length])

def encode_double(value):
	return struct.pack("<d", value)

def decode_double(data, base):
	return (base + 8, struct.unpack("<d", data[base: base + 8])[0])


ELEMENT_TYPES = {
		0x01 : "double",
		0x02 : "string",
		0x03 : "document",
		0x04 : "array",
		0x05 : "binary",
		0x08 : "boolean",
				0x09 : "UTCdatetime",
		0x0A : "none",
		0x10 : "int32",
		0x12 : "int64"
	}

def encode_double_element(name, value):
	return "\x01" + encode_cstring(name) + encode_double(value)

def decode_double_element(data, base):
	base, name = decode_cstring(data, base + 1)
	base, value = decode_double(data, base)
	return (base, name, value)

def encode_string_element(name, value):
	return "\x02" + encode_cstring(name) + encode_string(value)

def decode_string_element(data, base):
	base, name = decode_cstring(data, base + 1)
	base, value = decode_string(data, base)
	return (base, name, value)

def encode_value(name, value, buf, traversal_stack, generator_func):
	if isinstance(value, BSONCoding):
		buf.write(encode_object_element(name, value, traversal_stack,
			generator_func))
	elif isinstance(value, float):
		buf.write(encode_double_element(name, value))
	elif isinstance(value, unicode):
		buf.write(encode_string_element(name, value))
	elif isinstance(value, dict):
		buf.write(encode_document_element(name, value,
			traversal_stack, generator_func))
	elif isinstance(value, list) or isinstance(value, tuple):
		buf.write(encode_array_element(name, value,
			traversal_stack, generator_func))
	elif isinstance(value, str):
		buf.write(encode_binary_element(name, value))
	elif isinstance(value, bool):
		buf.write(encode_boolean_element(name, value))
	elif isinstance(value, datetime):
		buf.write(encode_UTCdatetime_element(name, value))
	elif value is None:
		buf.write(encode_none_element(name, value))
	elif isinstance(value, int):
		if value < -0x80000000 or value > 0x7fffffff:
			buf.write(encode_int64_element(name, value))
		else:
			buf.write(encode_int32_element(name, value))
	elif isinstance(value, long):
		buf.write(encode_int64_element(name, value))

def encode_document(obj, traversal_stack,
		traversal_parent = None,
		generator_func = None):
	buf = cStringIO.StringIO()
	key_iter = obj.iterkeys()
	if generator_func is not None:
		key_iter = generator_func(obj, traversal_stack)
	for name in key_iter:
		value = obj[name]
		traversal_stack.append(TraversalStep(traversal_parent or obj, name))
		encode_value(name, value, buf, traversal_stack, generator_func)
		traversal_stack.pop()
	e_list = buf.getvalue()
	e_list_length = len(e_list)
	return struct.pack("<i%dsb" % (e_list_length,), e_list_length + 4 + 1,
			e_list, 0)

def encode_array(array, traversal_stack,
		traversal_parent = None,
		generator_func = None):
	buf = cStringIO.StringIO()
	for i in xrange(0, len(array)):
		value = array[i]
		traversal_stack.append(TraversalStep(traversal_parent or array, i))
		encode_value(unicode(i), value, buf, traversal_stack, generator_func)
		traversal_stack.pop()
	e_list = buf.getvalue()
	e_list_length = len(e_list)
	return struct.pack("<i%dsb" % (e_list_length,), e_list_length + 4 + 1,
			e_list, 0)

def decode_element(data, base):
	element_type = struct.unpack("<b", data[base:base + 1])[0]
	element_description = ELEMENT_TYPES[element_type]
	decode_func = globals()["decode_" + element_description + "_element"]
	return decode_func(data, base)

def decode_document(data, base):
	length = struct.unpack("<i", data[base:base + 4])[0]
	end_point = base + length
	base += 4
	retval = {}
	while base < end_point - 1:
		base, name, value = decode_element(data, base)
		retval[name] = value
	if "$$__CLASS_NAME__$$" in retval:
		retval = decode_object(retval)
	return (end_point, retval)

def encode_document_element(name, value, traversal_stack, generator_func):
	return "\x03" + encode_cstring(name) + \
			encode_document(value, traversal_stack,
					generator_func = generator_func)

def decode_document_element(data, base):
	base, name = decode_cstring(data, base + 1)
	base, value = decode_document(data, base)
	return (base, name, value)

def encode_array_element(name, value, traversal_stack, generator_func):
	return "\x04" + encode_cstring(name) + \
			encode_array(value, traversal_stack, generator_func = generator_func)

def decode_array_element(data, base):
	base, name = decode_cstring(data, base + 1)
	base, value = decode_document(data, base)
	retval = []
	try:
		i = 0
		while True:
			retval.append(value[unicode(i)])
			i += 1
	except KeyError:
		pass
	return (base, name, retval)

def encode_binary_element(name, value):
	return "\x05" + encode_cstring(name) + encode_binary(value)

def decode_binary_element(data, base):
	base, name = decode_cstring(data, base + 1)
	base, value = decode_binary(data, base)
	return (base, name, value)

def encode_boolean_element(name, value):
	return "\x08" + encode_cstring(name) + struct.pack("<b", value)

def decode_boolean_element(data, base):
	base, name = decode_cstring(data, base + 1)
	value = not not struct.unpack("<b", data[base:base + 1])[0]
	return (base + 1, name, value)

def encode_UTCdatetime_element(name, value):
	if value.tzinfo is None:
		warnings.warn(MissingTimezoneWarning(), None, 4)
	value = int(round(calendar.timegm(value.utctimetuple()) * 1000 +
		(value.microsecond / 1000.0)))
	return "\x09" + encode_cstring(name) + struct.pack("<q", value)

def decode_UTCdatetime_element(data, base):
	base, name = decode_cstring(data, base + 1)
	value = datetime.fromtimestamp(struct.unpack("<q",
		data[base:base + 8])[0] / 1000.0, pytz.utc)
	return (base + 8, name, value)

def encode_none_element(name, value):
	return "\x0a" + encode_cstring(name)

def decode_none_element(data, base):
	base, name = decode_cstring(data, base + 1)
	return (base, name, None)

def encode_int32_element(name, value):
	return "\x10" + encode_cstring(name) + struct.pack("<i", value)

def decode_int32_element(data, base):
	base, name = decode_cstring(data, base + 1)
	value = struct.unpack("<i", data[base:base + 4])[0]
	return (base + 4, name, value)

def encode_int64_element(name, value):
	return "\x12" + encode_cstring(name) + struct.pack("<q", value)

def decode_int64_element(data, base):
	base, name = decode_cstring(data, base + 1)
	value = struct.unpack("<q", data[base:base + 8])[0]
	return (base + 8, name, value)
# }}}

args = list(args)
		if len(args) < 1:
			args.append("Input datetime object has no tzinfo, assuming UTC.")
		super(MissingTimezoneWarning, self).__init__(*args)