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
from networkers import NetworkReader,NetworkWriter


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:
			indices=[]
			for i in range(256):
				indices.append(nodeStack.pop())
				if len(nodeStack)==0: break
			indices.sort()
			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))

			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]:
					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()

		return sorted(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)

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

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