import hashlib
import socket
import logging as log

from hashtree import HashTree
from networkers import NetworkReader,NetworkWriter
import config as conf


class Connection:
	def __init__(self,serverSocket):
		self.socket, address = serverSocket.accept()
		log.info('Connected by {0}'.format(address))
		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.close()


class Server:
	def __init__(self,filename,treeFile=""):
		self.filename=filename
		self._treeFile=treeFile

		if treeFile:
			self.tree=HashTree.load(treeFile)
		else:
			self.tree=HashTree.fromFile(filename)

		self._newLeaves=dict()
		self.BLOCK_SIZE=self.tree.BLOCK_SIZE

		self.ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		self.ss.bind(("",conf.port))
		self.ss.listen(1)

		self._lastWrite=-1
		self.dataFile=None

	def serve(self):
		while True:
			with Connection(self.ss) as (incoming,outcoming):
				try:
					while True:
						if not self._serveOne(incoming,outcoming): return
				except AssertionError:
					continue

	def _serveOne(self,incoming,outcoming):
		jsonData,binData=incoming.readMsg()

		if jsonData["command"]=="init":
			assert jsonData["blockSize"]==self.BLOCK_SIZE
			assert jsonData["blockCount"]==self.tree.leafCount
			outcoming.writeMsg({"command": "ack"})

		elif jsonData["command"]=="req":
			outcoming.writeMsg(*self._requestHash(jsonData))

		elif jsonData["command"]=="send" and jsonData["dataType"]=="data":
			outcoming.writeMsg(*self._receiveData(jsonData,binData))

		elif jsonData["command"]=="end":
			self._finalize()
			return False

		else:
			assert False, jsonData["command"]

		return True

	def _requestHash(self,jsonData):
		log.info("received request for node #{0}".format(jsonData["index"]))
		assert jsonData["index"]<len(self.tree.store)
		nodeHash=self.tree.store[jsonData["index"]]

		jsonResponse={"command":"send", "index":jsonData["index"], "dataType":"hash"}
		binResponse=nodeHash

		return (jsonResponse,binResponse)

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

		if not self.dataFile:
			self.dataFile=open(self.filename,mode="rb+")
		i=jsonData["index"]
		if self._lastWrite+1!=i:
			self.dataFile.seek(i*self.BLOCK_SIZE)
		self.dataFile.write(binData)
		self._lastWrite=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...")
		if self.dataFile:
			self.dataFile.close()
		if self._treeFile:
			log.info("updating hash tree...")
			for (k,v) in self._newLeaves:
				self.tree.updateLeaf(k,v)
			self.tree.save(self._treeFile)
		log.info("done")