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 rcp_node import RCP_Quorum
from collections import defaultdict 

class QP_Quorum(RCP_Quorum):
    def __init__(self, quorumID, documentSize, numNodes, numItems=0, table=[]):
        RCP_Quorum.__init__(self, quorumID, documentSize, numNodes, numItems, table)
        self.fingerTableRangeAccesses = [defaultdict(lambda: 0) for i in range(self.numNodes)]
        self.fingerTableAccesses = [defaultdict(lambda: 0) for i in range(self.numNodes)]

    def get_finger_table_range_accesses(self, whichNode):
        return self.fingerTableRangeAccesses[whichNode]

    def get_finger_table_accesses(self, whichNode):
        return self.fingerTableAccesses[whichNode]

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

    def get_finger_table_ranges(self, whichNode, numKeys, numSignatures):
        numEntries = len(self.table.items())
        retval = (list(self.table.keys()), numEntries)

        # 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 finger table
        sizeOfResponse = SIZE_OF_HASH * numEntries
        # Then the entrywise encrypted finger table routing information
        sizeOfResponse += (SIZE_OF_IP_ADDRESS * self.numNodes + SIZE_OF_KEY) * numEntries
        # Then the OT prime values
        sizeOfResponse += 2 * SIZE_OF_OT_VALUE * numEntries
        # 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.fingerTableRangeAccesses[whichNode][numEntries] += 1

        return retval

    def OT_get_finger_table_val(self, whichNode, index):
        numEntries = len(self.table.items())
        retval = self.table[index]

        # OT prime value to be able to open one of the entries in the finger table
        sizeOfRequest = SIZE_OF_OT_VALUE
        # And a signature to confirm it's valid
        sizeOfRequest += SIZE_OF_SIGNATURE

        # The final OT prime value
        sizeOfResponse = SIZE_OF_OT_VALUE
        # And a signature to boot
        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.fingerTableAccesses[whichNode][numEntries] += 1

        return retval

# TODO: Add unit tests for size calculations
# TODO: Add unit test to make sure finger_table_ranges/OT_val is still firing correctly?
if __name__ == "__main__":
    SIZE_OF_DOCUMENTS_IN_TEST = 1024
    NUM_NODES_PER_QUORUM_IN_TEST = 10
    test = QP_Quorum(0, SIZE_OF_DOCUMENTS_IN_TEST, NUM_NODES_PER_QUORUM_IN_TEST)

    randomNode = 0
    [test.insert() for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    [test.retrieve(randomNode) for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    print("Insert and retrieval on node 0 fires correctly.")

    [test.retrieve(x) for x in range(NUM_NODES_PER_QUORUM_IN_TEST)]
    print("Retrieval fires on all nodes correctly.")