import hashlib
import os
import collections


class HashTree:
  HASH_LEN=16 # bytes
  BLOCK_SIZE=4096 # bytes
  
  ## Prepares a tree containing leafCount leaves.
  def __init__(self,leafCount):
    self.store=[None]*(leafCount*2-1)
    self.leafStart=leafCount-1
    self.index=self.leafStart
    
  @classmethod
  def fromFile(cls,fd):
    stat=os.fstat(fd.fileno())
    size=stat.st_size # !! symlinks
    leafCount=(size-1)//HashTree.BLOCK_SIZE+1 # number of leaf blocks
    res=cls(leafCount)
    
    for i in range(leafCount):
      data=fd.read(HashTree.BLOCK_SIZE)
      res.insertLeaf(hashlib.sha256(data).digest()[HashTree.HASH_LEN:])
    res.buildTree()
    
    return res
    
    
  ## Inserts a leaf at the first empty position.
  #  
  #  Useful and used only during the tree construction.
  def insertLeaf(self,h):
    self.store[self.index]=h
    self.index+=1
    
  ## Updates a hash stored in the leaf.
  def updateLeaf(self,index,h):
    if index<self.leafStart: raise IndexError()
    
    self.store[index]=h
    self.updateNode((index-1)//2)
  
  ## Updates the node at index and all its ancestors.
  def updateNode(self,index):
    while index>=0:
      self.store[index]=hashlib.sha256(self.store[index*2+1]+self.store[index*2+2]).digest()[HashTree.HASH_LEN:]
      index=(index-1)//2
      
  ## Fast construction of the tree over the leaves. O(n).
  def buildTree(self):
    for i in range(self.leafStart-1,-1,-1):
      self.store[i]=hashlib.sha256(self.store[i*2+1]+self.store[i*2+2]).digest()[HashTree.HASH_LEN:]


if __name__=="__main__":
  f1=HashTree.fromFile(open("serverFile.txt",mode='rb'))
  f2=HashTree.fromFile(open("clientFile.txt",mode='rb'))

  for i,(h1,h2) in enumerate(zip(f1.store,f2.store)):
    print("{0:2}".format(i),h1.hex(),h2.hex(),h1==h2)