class Connection:

	def __init__(self):

		sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM)

		sslContext=ssl.create_default_context(cafile=conf.peers)

		sslContext.check_hostname=False

		sslContext.load_cert_chain(conf.certfile,conf.keyfile)

		self._socket=sslContext.wrap_socket(sock)

		self._socket.connect((conf.hosts[0], conf.port))

		fr=self._socket.makefile(mode="rb")

		fw=self._socket.makefile(mode="wb")

		self.incoming=NetworkReader(fr)

		self.outcoming=NetworkWriter(fw)

	def __enter__(self):

		return self.incoming,self.outcoming

	def __exit__(self, exc_type, exc_val, exc_tb):

		self._socket.shutdown(socket.SHUT_RDWR)

		self._socket.close()

class Client:

	def __init__(self,filename):

		self._incoming=None

		self._outcoming=None

		self._filename=filename

		print(datetime.now(), "initializing...")

		self._localTree=HashTree.fromFile(self._filename)

	def negotiate(self):

		localTree=self._localTree

		blocksToTransfer=[]

		nodeStack=collections.deque([0]) # root

		# initialize session

		jsonData={"command":"init", "blockSize":localTree.BLOCK_SIZE, "blockCount":localTree.leafCount, "version":conf.version}

		self._outcoming.writeMsg(jsonData)

		jsonData,binData=self._incoming.readMsg()

		assert jsonData["command"]=="ack"

		# determine which blocks to send

		print(datetime.now(), "negotiating:")

		progress=Progress(localTree.leafCount)

		while len(nodeStack)>0:

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["dataType"]=="hash"

				if 2*i+3<len(localTree.store): # inner node

					nodeStack.append(2*i+2)

					nodeStack.append(2*i+1)

				else:

					blocksToTransfer.append(i-localTree.leafStart) # leaf

					progress.p(i-localTree.leafStart)

		progress.done()

	def sendData(self,blocksToTransfer):

		log.info(blocksToTransfer)

		dataFile=open(self._filename, mode="rb")

		i1=-1

		print(datetime.now(), "sending data:")

		progress=Progress(len(blocksToTransfer))

		for (k,i2) in enumerate(blocksToTransfer):

			jsonData={"command":"send", "index":i2, "dataType":"data"}

			if i1+1!=i2:

				dataFile.seek(i2*HashTree.BLOCK_SIZE)

			binData=dataFile.read(HashTree.BLOCK_SIZE)

			log.info("block #{0}: {1}...{2}".format(i2,binData[:5],binData[-5:]))

			self._outcoming.writeMsg(jsonData,binData)

			stats.logTransferredBlock()

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["command"]=="ack" and jsonData["index"]==i2, jsonData

			i1=i2

			progress.p(k)

		progress.done()

		self._outcoming.writeMsg({"command":"end"})

		log.info("closing session...")

		dataFile.close()

	def pullData(self,blocksToTransfer):

		log.info(blocksToTransfer)

		dataFile=open(self._filename, mode="rb+")

		i1=-1

		print(datetime.now(), "receiving data:")

		progress=Progress(len(blocksToTransfer))

		for (k,i2) in enumerate(blocksToTransfer):

			self._outcoming.writeMsg({"command":"req", "index":i2, "dataType":"data"})

			jsonData,binData=self._incoming.readMsg()

			assert jsonData["command"]=="send" and jsonData["index"]==i2 and jsonData["dataType"]=="data", jsonData

			if i1+1!=i2:

				dataFile.seek(i2*HashTree.BLOCK_SIZE)

			dataFile.write(binData)

			log.info("block #{0}: {1}...{2}".format(i2,binData[:5],binData[-5:]))

			stats.logTransferredBlock()

		self._lastIndex=-1

		self._dataFileHandle=None

	@staticmethod

	def run(*args):

		s=Server(*args)

		s.serve()

	@property

	def _dataFile(self):

		if not self._dataFileHandle:

			self._dataFileHandle=open(self._filename, mode="rb+")

		return self._dataFileHandle

	def serve(self):

		try:

			while self._serveOne(): pass

		except AssertionError as e:

			log.warning(e)

	def _serveOne(self):

		jsonData,binData=self._incoming.readMsg()

		if jsonData["command"]=="init":

			assert jsonData["blockSize"]==self.BLOCK_SIZE

			assert jsonData["blockCount"]==self._tree.leafCount

			self._outcoming.writeMsg({"command": "ack"})

		elif jsonData["command"]=="req":

			if jsonData["dataType"]=="data":

				self._outcoming.writeMsg(*self._requestData(jsonData["index"]))

			else:

				self._outcoming.writeMsg(*self._requestHash(jsonData["index"]))

		elif jsonData["command"]=="send" and jsonData["dataType"]=="data":

			self._outcoming.writeMsg(*self._receiveData(jsonData,binData))

		elif jsonData["command"]=="end":

			self._finalize()

			self._locked=False

			return False

		else:

			assert False, jsonData["command"]

		return True

		return (jsonResponse,binResponse)

	def _requestData(self,index):

		log.info("received request for data block #{0}".format(index))

		jsonResponse={"command":"send", "index":index, "dataType":"data"}

		if self._lastIndex+1!=index:

			self._dataFile.seek(index*self.BLOCK_SIZE)

		binResponse=self._dataFile.read(self.BLOCK_SIZE)

		return (jsonResponse,binResponse)

	def _receiveData(self,jsonData,binData):

		log.info("received data block #{0}: {1}...{2}".format(jsonData["index"],binData[:5],binData[-5:]))

		i=jsonData["index"]

		if self._lastIndex+1!=i:

			self._dataFile.seek(i*self.BLOCK_SIZE)

		self._dataFile.write(binData)

		self._lastIndex=i

		if self._treeFile:

			self._newLeaves[i+self._tree.leafStart]=hashlib.sha256(binData).digest()[HashTree.HASH_LEN:]

		return ({"command": "ack", "index": i},)

	def _finalize(self):

		log.info("closing session...")

		self._dataFile.close()

		self._dataFileHandle=None

		if self._treeFile:

			log.info("updating hash tree...")

			for (k,v) in self._newLeaves.items():

				self._tree.updateLeaf(k, v)

			self._tree.save(self._treeFile)

		log.info("done")

class Stats:

	received=0

	sent=0

	exchangedNodes=0

	transferredBlocks=0

def logReceived(data):

	Stats.received+=len(data)

def logSent(data):

	Stats.sent+=len(data)

def logTransferredBlock():

	Stats.transferredBlocks+=1

def report():

	return """received {r}B

sent {s}B

exchanged {nodes} hash tree nodes

transferred {blocks} blocks""".format(r=Stats.received, s=Stats.sent, nodes=Stats.exchangedNodes, blocks=Stats.transferredBlocks)