import collections
import socket
import ssl
import logging as log
from datetime import datetime

import config as conf
import stats
from util import Progress
from hashtree import HashTree,hashBlock
from netnode import BaseConnection,NetNode,FailedConnection


class Connection(BaseConnection):
	def __init__(self,host,port):
		super().__init__()
		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)

		try:
			self._socket.connect((host,port))
		except ConnectionRefusedError as e:
			log.exception(e)
			print("Couldn't connect to {0}:{1}".format(host,port))
			raise FailedConnection()
		except ssl.SSLError as e:
			log.exception(e)
			print("Error creating SSL connection to {0}:{1}".format(host,port))
			raise FailedConnection()

		self.createNetworkers()
		print("Connected to {0}".format(host))


class Client(NetNode):
	def __init__(self,filename,treeFile=""):
		print(datetime.now(), "initializing...")
		super().__init__(filename,treeFile)

	## Asks server for node hashes to determine which are to be transferred.
	#
	# Uses a binary HashTree, where item at k is hash of items at 2k+1, 2k+2.
	#
	# Requests nodes in order of a batch DFS. Needs stack of size O(treeDepth*batchSize). Nodes in each tree level are accessed in order.
	def negotiate(self):
		localTree=self._tree
		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:
			indices=[]
			for i in range(conf.batchSize):
				indices.append(nodeStack.pop())
				if len(nodeStack)==0: break
			self._outcoming.writeMsg({"command":"req", "index":indices, "dataType":"hash"})

			jsonData,binData=self._incoming.readMsg()
			assert jsonData["index"]==indices
			assert jsonData["dataType"]=="hash"
			stats.logExchangedNode(len(indices))

			frontier=[]
			for (j,i) in enumerate(indices):
				(j1,j2)=[HashTree.HASH_LEN*ji for ji in (j,j+1)]
				if localTree.store[i]!=binData[j1:j2]:
					# ie. 0-6 nodes, 7-14 leaves. 2*6+2<15
					if 2*i+2<len(localTree.store): # inner node
						frontier.append(2*i+1)
						frontier.append(2*i+2)
					else:
						blocksToTransfer.append(i-localTree.leafStart) # leaf
						progress.p(i-localTree.leafStart)
			nodeStack.extend(reversed(frontier))
		progress.done()

		size=stats.formatBytes(len(blocksToTransfer)*self._tree.BLOCK_SIZE)
		print(datetime.now(), "{0} to transfer".format(size))

		return blocksToTransfer

	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)

			if self._treeFile:
				self._newLeaves[i2+self._tree.leafStart]=hashBlock(binData)

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

			stats.logTransferredBlock()
			i1=i2
			progress.p(k)
		progress.done()

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

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

		if self._treeFile:
			self._updateTree()

	def setConnection(self,connection):
		(self._incoming,self._outcoming)=connection