from dht_common import SIZE_OF_HASH, SIZE_OF_IP_ADDRESS, SIZE_OF_OT_VALUE, SIZE_OF_KEY, SIZE_OF_SIGNATURE, SIZE_OF_TIMESTAMP
from qp_node import QP_Quorum
from collections import defaultdict 

class QPLastHop_Quorum(QP_Quorum):
    def __init__(self, quorumID, documentSize, numNodes, numItems=0, table=[]):
        QP_Quorum.__init__(self, quorumID, documentSize, numNodes, numItems, table)
        self.databaseAccesses = [defaultdict(lambda: 0) for i in range(self.numNodes)]
        
    def get_database_accesses(self, whichNode):
        return self.databaseAccesses[whichNode]

    def get_final_table_ranges(self, whichNode, numKeys, numSignatures):
        # Asker's ID
        sizeOfRequest = SIZE_OF_HASH
        # timestamp
        sizeOfRequest += SIZE_OF_TIMESTAMP
        # keys in request
        sizeOfRequest += SIZE_OF_KEY * numKeys
        # signatures in request
        sizeOfRequest += SIZE_OF_SIGNATURE * numSignatures

        # The set of hashes in the item store
        sizeOfResponse = SIZE_OF_HASH * self.numItems
        # Then the entrywise encrypted item store
        sizeOfResponse += self.documentSize * self.numItems
        # Then the OT prime values
        sizeOfResponse += 2 * SIZE_OF_OT_VALUE * self.numItems
        # Then, finally, a signature to tie it together
        sizeOfResponse += SIZE_OF_SIGNATURE

        self.nodeNumRounds[whichNode] += 1
        self.nodeNumMessagesSent[whichNode] += 1
        self.nodeNumMessagesRecv[whichNode] += 1
        self.nodeNumBytesSent[whichNode] += sizeOfResponse
        self.nodeNumBytesRecv[whichNode] += sizeOfRequest

        self.databaseAccesses[whichNode][self.numItems] += 1

        return self.numItems

    # This shouldn't be used, just here to make sure you don't try the RCP_Quorum function it overrides
    def retrieve(self):
        return None

    def OT_retrieve(self, whichNode, numKeys, numSignatures):
        # Asker's ID
        sizeOfRequest = SIZE_OF_HASH
        # timestamp
        sizeOfRequest += SIZE_OF_TIMESTAMP
        # keys in request
        sizeOfRequest += SIZE_OF_KEY * numKeys
        # signatures in request
        sizeOfRequest += SIZE_OF_SIGNATURE * numSignatures
        # actual OT crypto usage
        sizeOfRequest += SIZE_OF_OT_VALUE
        # signature on whole thing
        sizeOfRequest += SIZE_OF_SIGNATURE

        sizeOfResponse = SIZE_OF_OT_VALUE + SIZE_OF_SIGNATURE

        self.nodeNumRounds[whichNode] += 1
        self.nodeNumMessagesSent[whichNode] += 1
        self.nodeNumMessagesRecv[whichNode] += 1
        self.nodeNumBytesSent[whichNode] += sizeOfResponse
        self.nodeNumBytesRecv[whichNode] += sizeOfRequest

# TODO: Add unit tests for size calculations
if __name__ == "__main__":
    SIZE_OF_DOCUMENTS_IN_TEST = 1024
    NUM_NODES_PER_QUORUM_IN_TEST = 10
    test = QPLastHop_Quorum(0, SIZE_OF_DOCUMENTS_IN_TEST, NUM_NODES_PER_QUORUM_IN_TEST)
    
    [test.insert() for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    
    [test.get_final_table_ranges(x, 0, 0) for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    print("Getting final table ranges fires on all nodes correctly.")

    [test.OT_retrieve(x) for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    print("OT retrieval fires on all nodes correctly.")