import struct
import types
import cStringIO
import string

class IndexedPage:
	#-- pageType,pageLength,directoryLength
	PAGE_HEADER_STRUCTURE="!III"
	PAGE_HEADER_STRUCTURE_length=struct.calcsize(PAGE_HEADER_STRUCTURE)
	DIRECTORY_STRUCTURE="!IIB"
	DIRECTORY_STRUCTURE_length=struct.calcsize(DIRECTORY_STRUCTURE)

	#-- kwargs:
	#-- source=<opened file>
	#-- data=<datablock>
	def __init__(self,kwdict=None,**kwargs):
		self.storage="memory"
		self.directory={}
		self.cache={}
		self.pageType=-1

		self.source=None
		self.sourceOffset=0
		self.baseAddress=0

		if kwdict is not None:
			kwargs=kwdict

		if kwargs.has_key('source'):
			self.storage="disk"
			self.source=kwargs['source']

		elif kwargs.has_key('data'):
			#-- I'll use data to construct a temporary memory file.
			self.source=cStringIO.StringIO(kwargs['data'])
		else:
			pass

		if self.source is not None:
			#-- The offset of this index page relative to the file.
			self.sourceOffset=self.source.tell()

			#-- Read the page header.
			self.pageType,pageLength,directoryLength=struct.unpack(IndexedPage.PAGE_HEADER_STRUCTURE,self.source.read(IndexedPage.PAGE_HEADER_STRUCTURE_length))			
			#print pageLength,directoryLength

			#-- Compute the base address of the entries (not the directory).
			self.baseAddress=self.sourceOffset+IndexedPage.PAGE_HEADER_STRUCTURE_length+directoryLength

			#-- Read the directory, and parse it.
			directoryData=self.source.read(directoryLength)
			self.parseDirectory(directoryData)
	
			#-- Check for aggressive caching.
			if kwargs.has_key('cache'):
				#-- Yep, load the entire page into memory.
				self.cache=self.readAll()
				self.storage="memory"
			elif kwargs.has_key('data'):
				self.cache=self.readAll()
				self.source.close()
			else:
				#-- Nope.  Position the file pointer at the end of the page.
				self.source.seek(self.sourceOffset+pageLength)

	#-- Returns a list of keys that are in the index.
	def keys(self):
		if self.storage=="memory":
			return self.cache.keys()
		if self.storage=="disk":
			return self.directory.keys()

	def has_key(self,key):
		if self.storage=="memory":
			return self.cache.has_key(key)
		if self.storage=="disk":
			return self.directory.has_key(key)

	def __getitem__(self,indexTerm):
		if self.storage=="disk":
			return self.readEntry(indexTerm)
		else:
			return self.cache[indexTerm]

	def __setitem__(self,index,value):
		if self.storage=="disk":
			raise Exception("IndexedPage is on disk and is therefore read only")
		self.cache[index]=value

	def readEntry(self,indexTerm):
		entryOffset,entryLength=self.directory[indexTerm]
		self.source.seek(self.baseAddress+entryOffset)
		return self.unserializeEntry(self.source.read(entryLength))

	def readAll(self):
		cache={}
		for indexTerm in self.directory.keys():
			cache[indexTerm]=self.readEntry(indexTerm)
		return cache

	def parseDirectory(self,directoryData):
		i=0
		while i<len(directoryData):
			indexHeader=directoryData[i:i+IndexedPage.DIRECTORY_STRUCTURE_length]
			entryOffset,entryLength,indexLength=struct.unpack(IndexedPage.DIRECTORY_STRUCTURE,indexHeader)
			indexTerm=self.unserializeIndexTerm(directoryData[i+IndexedPage.DIRECTORY_STRUCTURE_length:i+IndexedPage.DIRECTORY_STRUCTURE_length+indexLength])
			self.directory[indexTerm]=(entryOffset,entryLength)
			i=i+IndexedPage.DIRECTORY_STRUCTURE_length+indexLength

	def unserializeIndexTerm(self,serializedIndexTerm):
		return serializedIndexTerm
	
	def serializeIndexTerm(self,indexTerm):
		return indexTerm

	def serializeEntry(self,entry):
		return entry

	def unserializeEntry(self,entryData):
		return entryData
				
	def __str__(self):
		if self.storage=="memory":
			cachedEntries=self.cache
		if self.storage=="disk":
			cachedEntries=self.readAll()

		#-- Serialize the entries.
		workingDirectory={}
		entryList=[]
		entryOffset=0
		for indexTerm in cachedEntries.keys():
			cachedEntry=self.serializeEntry(cachedEntries[indexTerm])
			entryList.append(cachedEntry)
			entryLength=len(cachedEntry)
			workingDirectory[indexTerm]=(entryOffset,entryLength)
			entryOffset=entryOffset+entryLength
		serializedEntries=treeReduce(lambda x,y:x+y,entryList)
		
		#-- Serialize the directory.
		directoryList=[]
		sortedIndexTerms=workingDirectory.keys()
		sortedIndexTerms.sort()
		for indexTerm in sortedIndexTerms:
			entryOffset,entryLength=workingDirectory[indexTerm]
			serializedIndexTerm=self.serializeIndexTerm(indexTerm)
			serializedIndexTermLength=len(serializedIndexTerm)
			directoryList.append(struct.pack(IndexedPage.DIRECTORY_STRUCTURE,entryOffset,entryLength,serializedIndexTermLength)+serializedIndexTerm)
		serializedDirectory=treeReduce(lambda x,y:x+y,directoryList)
		
		#-- Now, compute the page header.
		serializedDirectoryLength=len(serializedDirectory)
		serializedEntriesLength=len(serializedEntries)
		pageLength=IndexedPage.PAGE_HEADER_STRUCTURE_length+serializedDirectoryLength+serializedEntriesLength
		pageHeader=struct.pack(IndexedPage.PAGE_HEADER_STRUCTURE,self.pageType,pageLength,serializedDirectoryLength)

		return "%s%s%s"%(pageHeader,serializedDirectory,serializedEntries)


class IndexSet:
	def __init__(self,pageTypeNumber,pageTypeClass):
		self.pageTypeNumber=pageTypeNumber
		self.pageTypeClass=pageTypeClass
		self.indexPages=[]

	def keys(self):
		Keys={}
		for index in self.indexPages:
			subsetKeys=index.keys()
			for subsetKey in subsetKeys:
				Keys[subsetKey]=None
		return Keys.keys()
	
	def has_key(self,key):
		for index in self.indexPages:
			if index.has_key(key):
				return 1
		return 0

	def __getitem__(self,indexTerm):
		superHit=[]
		for index in self.indexPages:
			if index.has_key(indexTerm):
				hit=index[indexTerm]
				if type(hit)==types.ListType:
					superHit=superHit+hit
				else:
					superHit.append(hit)					
		if len(superHit)==0:
			return None
		elif len(superHit)==1:
			return superHit[0]
		else:
			return superHit
	
	def addIndex(self,index):
		self.indexPages.append(index)	

	def __str__(self):
		serializedIndeces=[]
		for index in self.indexPages:
			serializedIndeces.append('%s'%index)
		return treeReduce(lambda x,y:x+y,serializedIndeces)

class Index:
	def __init__(self,pageTypes,sourceFile=None):
		self.sourceFile=sourceFile
		self.pageTypes=pageTypes
		self.indexSets={}
		for pageType in self.pageTypes.keys():
			self.indexSets[pageType]=IndexSet(pageType,self.pageTypes[pageType])

		if self.sourceFile:
			pageHeaderData=self.sourceFile.read(IndexedPage.PAGE_HEADER_STRUCTURE_length)
			while len(pageHeaderData):
				#-- Read the page header.
				pageType,pageLength,directoryLength=struct.unpack(IndexedPage.PAGE_HEADER_STRUCTURE,pageHeaderData)
				self.sourceFile.seek(-IndexedPage.PAGE_HEADER_STRUCTURE_length,1)
				#-- Check to see if this is an interesting page.
				if self.pageTypes.has_key(pageType):
					#-- Yes it is!  Load the page.
					self.indexSets[pageType].addIndex(pageTypes[pageType](source=self.sourceFile))
				else:
					#-- No it's not.  Skip this page.
					self.sourceFile.seek(pageLength,1)
				pageHeaderData=self.sourceFile.read(IndexedPage.PAGE_HEADER_STRUCTURE_length)

	def __str__(self):
		serializedIndexSets=[]
		for pageType in self.indexSets.keys():
			print "Serializing page type: ",pageType
			serializedIndexSets.append('%s'%self.indexSets[pageType])
		return treeReduce(lambda x,y:x+y,serializedIndexSets)

	def __getitem__(self,pageType):
		return self.indexSets[pageType]

	def addIndex(self,pageType,index):
		self.indexSets[pageType].addIndex(index)

#-- f should be of the form z=f(x,y) where x,y,z are all singletons.
def treeReduce(f,list):
	l=list
	while len(l)>1:
		size=len(l)
		l2=[]
		for i in range(size/2):
			l2.append(f(l[i*2],l[i*2+1]))
		if size%2!=0:
			l2.append(l[size-1])
		l=l2
	return l[0]