from hashtree import HashTree
import collections
import socket
import sys

import config as conf
from networkers import NetworkReader,NetworkWriter


filename=sys.argv[1]


class Connection:
	def __init__(self):
		self.socket=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
		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.close()


def negotiate():
	localTree=HashTree.fromFile(open(filename,mode="rb"))
	blocksToTransfer=[]
	nodeStack=collections.deque([0]) # root

	# initialize session
	with Connection() as (incoming,outcoming):
		jsonData={"command":"init", "blockSize":localTree.BLOCK_SIZE, "blockCount":localTree.leafCount, "version":conf.version}
		outcoming.writeMsg(jsonData)

	# determine which blocks to send
	while len(nodeStack)>0:
		with Connection() as (incoming,outcoming):
			i=nodeStack.pop()
			outcoming.writeMsg({"command":"req", "index":i})

			jsonData,binData=incoming.readMsg()
			assert jsonData["index"]==i
			assert jsonData["dataType"]=="hash"

			if localTree.store[i]!=binData:
				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

	return blocksToTransfer


def sendData(blocksToTransfer):
	print(blocksToTransfer)
	dataFile=open(filename,mode="rb")
	i1=-1

	for i2 in blocksToTransfer:
		with Connection() as (incoming,outcoming):
			jsonData={"command":"send", "index":i2, "dataType":"data"}
			if i1+1!=i2:
				dataFile.seek(i2*HashTree.BLOCK_SIZE)
			binData=dataFile.read(HashTree.BLOCK_SIZE)

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

			outcoming.writeMsg(jsonData,binData)
			i1=i2

	with Connection() as (incoming,outcoming):
		outcoming.writeMsg({"command":"end"})

	print("closing...")
	dataFile.close()


if __name__=="__main__":
	blocksToTransfer=negotiate()
	sendData(blocksToTransfer)

	sys.exit(0)