import hashlib
import socket
import ssl
import multiprocessing
import logging as log

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


class Connection:
	def __init__(self,serverSocket,sslContext):
		sock, address = serverSocket.accept()
		self._socket=sslContext.wrap_socket(sock,server_side=True)

		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.shutdown(socket.SHUT_RDWR)
		self._socket.close()


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

		self._ssl=ssl.create_default_context(ssl.Purpose.CLIENT_AUTH,cafile=conf.peers)
		self._ssl.verify_mode=ssl.CERT_REQUIRED
		self._ssl.load_cert_chain(conf.certfile,conf.keyfile)

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

	def serve(self):
		while True:
			with Connection(self._ss,self._ssl) as c:
				p=multiprocessing.Process(target=Server.run,args=(c,self._filename,self._treeFile))
				p.start()
				p.join()


class Server:
	def __init__(self,connection,filename,treeFile=""):
		(self._incoming,self._outcoming)=connection
		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._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

	def _requestHash(self,indices):
		log.info("received request for nodes #{0}".format(",".join(str(i) for i in indices)))
		assert all(i<len(self._tree.store) for i in indices)
		hashes=[self._tree.store[i] for i in indices]

		jsonResponse={"command":"send", "index":indices, "dataType":"hash"}
		binResponse=b"".join(hashes)

		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")