import os

from BitTorrent.bitfield import Bitfield
from BitTorrent.PiecePicker import PiecePicker
from random import randrange, shuffle, choice
from time import time, strftime

import os
#from os import open,close,O_CREAT,O_EXCL

# I just need to add an additional function to the original Bitfield
# This is not used for the time being
class BitFieldWrapper(Bitfield):
	def have(self, index):
		return self[index]>0

#Need to extend _SingleTorrent
#Given input file, piece size and the # of the half pieces
# Write only the second half
# should be called by _SingleTorrent's _finish() function
def truncatefile(filename, piecesize, partialpieces):
	tmpfile = filename + '.tmp'
	fin = open(filename, 'r')
	fout = open(tmpfile, 'w')
	fin.seek(piecesize*partialpieces)
	while 1: 
		buffer=fin.read(1024)
		if (len(buffer)<=0):
			break
		fout.write(buffer)
	fin.close()
	fout.close()
	os.unlink(filename)
	os.rename(tmpfile, filename)

def truncatefileN(filename, piecesize, offset, partialpieces):
	tmpfile = filename + '.tmp'
	fin = open(filename, 'r')
	fout = open(tmpfile, 'w')
	fin.seek(piecesize*offset)
	count = 0
	while count<partialpieces: 
		buffer=fin.read(piecesize)
		if (len(buffer)<=0):
			break
		fout.write(buffer)
		count += 1
	fin.close()
	fout.close()
	os.unlink(filename)
	os.rename(tmpfile, filename)

def logMessage(config, msg): 
	t = strftime('[%H:%M:%S] ')
	log = t + msg + "\n"
	if config.has_key('logfile') and len(config['logfile'])>0: 
		f = open(config['logfile'],'a')
		if f:
			f.write(log)
			f.close()
		else:
			print log
		
# Given config (including processid) and total # of pieces, 
# return offset and the number of pieces
def getOffsetNPiece(config, numpieces):
	if config.has_key('processid') and (config['processid']>0):
		if (config.has_key('processes')==0):
			raise ValueError('# of processes must be defined')
		processes = config['processes']
		if (config['processid'] < processes):
			offset = (config['processid']-1) * (numpieces/processes)
			partialpieces = (numpieces/processes)
		elif (config['processid'] == processes):
			offset = (config['processid']-1) * (numpieces/processes)
			partialpieces = numpieces - (config['processid']-1) * (numpieces/processes)
		else:
			raise ValueError('processid must be less than processes')
	else:
		offset = 0
		partialpieces = numpieces
	return [offset, partialpieces]



# numpieces are still the ones
# I need to ensure that when requested is called, the piece should be within 
# [offset, offset+numpieces/2)
# "complete" should also be within this range
# next only returns that range as well (but as long as we ensure
# scrambled only returns that range, then this can be guaranteed
# am_I_complete: I just need to get the half

class RangedPiecePicker(PiecePicker): 
	def __init__(self, numpieces, config):
		PiecePicker.__init__(self, numpieces, config) 
		self.config = config
		if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1):
			(offset, partialpieces) = getOffsetNPiece(config, numpieces)
			self.offset = offset
			self.partialpieces = partialpieces
			self.scrambled = range(offset, offset+partialpieces)
			shuffle(self.scrambled)
		self._requested = [0] * numpieces
		self._others = [0] * numpieces
	def am_I_complete(self): 
		if self.config.has_key('processid') and (self.config['processid']>0): 
			if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1): 
				return self.numgot == self.partialpieces 
			else: 
				return self.numgot + sum(self._others) == self.numpieces
		else:
			return PiecePicker.am_I_complete(self)
	def got_have(self, piece):
	#	logMessage(self.config, "Others have piece " + str(piece) + "\n") 
		if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1): 
			if ((piece>=self.offset) and (piece < self.offset + self.partialpieces)): 
				PiecePicker.got_have(self, piece) 
		else:
			PiecePicker.got_have(self, piece)
	def lost_have(self, piece): 
		if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1): 
			if ((piece>=self.offset) and (piece < self.offset + self.partialpieces)): 
				PiecePicker.lost_have(self, piece) 
		else:
			PiecePicker.lost_have(self, piece)
	#def next(self, havefunc, seed = False): 
	#	bests = None 
	#	bestnum = 2 ** 30 
	#	if seed: 
	#		s = self.seedstarted 
	#	else: 
	#		s = self.started 
	#	for i in s: 
	#		if havefunc(i): 
	#			if self.numinterests[i] < bestnum: 
	#				bests = [i] 
	#				bestnum = self.numinterests[i] 
	#			elif self.numinterests[i] == bestnum: 
	#				bests.append(i) 
	#	if bests: 
	#		return choice(bests) 
	#	if self.numgot < self.config['rarest_first_cutoff']: 
	#		for i in self.scrambled: 
	#			if havefunc(i): 
	#				return i 
	#		return None 
	#	for i in xrange(1, min(bestnum, len(self.interests))): 
	#		for j in self.interests[i]: 
	#			if havefunc(j): 
	#				return j 
	#	return None
	# Make sure we only return the pieces that in the range
	def saveStorageWraper(self, _storagewrapper):
		self._storagewrapper = _storagewrapper
	def next(self, havefunc, seed = False): 
		if self.config.has_key('processid') and (self.config['processid']>0):
			pass
		else:
			piece = PiecePicker.next(self, havefunc, seed)
			if (piece == None):
				#logMessage(self.config, "We have in started with interests: " + str(self.interests))
				pass
			return piece
		def newhavefunc(index):
			return havefunc(index) and self._others[index]==0
		maxtries = 50
		tries = 0
		while (tries<maxtries):
			#piece = PiecePicker.next(self, havefunc, seed)
			piece = PiecePicker.next(self, newhavefunc, seed)
			#logMessage(self.config, 'Piece picker returns: ' + str(piece))
			if (piece == None):
				return None
			# Now need to check if the piece is there or not
			if self._requested[piece]==1:
				return piece
			tries += 1
			if self.isPieceThere(piece):
				logMessage(self.config, 'Piece ' + str(piece) + ' returned but exists') 
			else: 
				# For fixed piece assignment only 
				if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1): 
					if ( (piece>=self.offset) and (piece < self.offset+self.partialpieces)):
						logMessage(self.config, 'Piece ' + str(piece) + ' to download') 
						return piece
				else:
					logMessage(self.config, 'Piece ' + str(piece) + ' to download') 
					return piece
		logMessage(self.config, 'Too bad we are here')
		return None
	# If piece is already there, return True; else create the piece empty and return False
	def isPieceThere(self, piece):
		if self.config.has_key('piece_filename'): 
			filename =  self.config['piece_filename'] + str(piece) 
		if os.access(filename, os.F_OK):
			# Piece exists already
			logMessage(self.config, 'Piece ' + str(piece) + ' already exists') 
			self._others[piece]=1
			self._storagewrapper.inactive_requests[piece]=0
			return True
		try:
			f = os.open(filename, os.O_CREAT|os.O_EXCL)
			os.close(f)
		except OSError:
			# Piece exists already
			logMessage(self.config, 'Piece ' + str(piece) + ' already exists') 
			self._others[piece]=1
			self._storagewrapper.inactive_requests[piece]=0
			return True
		else: 
			# Piece not exist and we've created an empty file successfully
			self._requested[piece]=1
			return False
	# This is not used! 
	def isPieceThereNoLock(self, piece):
		if self.config.has_key('piece_filename'): 
			filename =  self.config['piece_filename'] + str(piece) 
		if not os.path.exists(filename):
			# create an empty file
			file = open(filename, "w")
			file.close()
			self._requested[piece]=1
			return False
		else: 
			self._others[piece]=1
			self._storagewrapper.inactive_requests[piece]=0
			return True
	def bump(self, piece):
		if self.config.has_key('use_fixed_piece_assignment') and (self.config['use_fixed_piece_assignment']==1): 
			if ((piece>=self.offset) and (piece < self.offset + self.partialpieces)):
				PiecePicker.bump(self, piece)
		else:
			PiecePicker.bump(self, piece)
	def complete(self, piece):
		if self.have[piece]:
			logMessage(self.config, 'We have piece ' + str(piece) + ' already')
			self.have[piece] = False
		PiecePicker.complete(self, piece)


if __name__ == '__main__': 
	config = {'rarest_first_cutoff':4}
	numpieces = 10
	offset = 5
	picker = RangedPiecePicker(numpieces, offset, numpieces/2, config)
	l = chr(255) + chr(128+64)
	bits = BitFieldWrapper(numpieces,l)
	print bits.have(9)
	print bits.have(4)
	picker.requested(5)
	a = picker.next(bits.have)
	print a
	picker.requested(6)
	a = picker.next(bits.have)
	print a
	picker.complete(6)
	a = picker.next(bits.have)
	print a
	picker.requested(7)
	picker.requested(8)
	picker.requested(9)
	picker.complete(5)
	picker.complete(9)
	a = picker.next(bits.have)
	print a
	picker.complete(8)
	a = picker.next(bits.have)
	print a
	picker.complete(7)
	a = picker.am_I_complete()
	print a