from dht_common import compute_document_ID, MAX_ID
from base_node import Base_Node
import math

class DHT_Simulator(object):
    def __init__(self, nodeType, numGroups, documentSize, numNodes):
        if nodeType == Base_Node:
            self.nodes = [nodeType(self.index_to_owner_node(i, numGroups), documentSize) for i in range(numGroups)]
        else:
            self.nodes = [nodeType(self.index_to_owner_node(i, numGroups), documentSize, numNodes) for i in range(numGroups)]
        self.__init_finger_tables()

    # Nobody else needs to call this, but when we put together all the nodes they need their finger table values updated
    def __init_finger_tables(self):
        numNodes = len(self.nodes)
        numEntries = math.ceil(math.log(numNodes, 2))
        [[self.nodes[i].insert_relation(self.index_to_owner_node(x), x) for x in [(2**j + i) % numNodes for j in range(numEntries)]] for i in range(numNodes)]

    # convert an index (of its internal data structure; consider this analogous to routing information)
    # to which node it should belong to (note that IDs are SHA3 outputs)
    def index_to_owner_node(self, which, numNodes=None):
        if not numNodes:
            numNodes = len(self.nodes)
        fullID = math.floor(which * MAX_ID / numNodes)
        return (fullID).to_bytes(32, byteorder="big")

    # convert a node ID (SHA3 output) to which index in the internal data structure it should be
    def owner_node_to_index(self, nodeID):
        numNodes = len(self.nodes)
        v = int.from_bytes(nodeID, byteorder='big')
        return math.floor(v * numNodes / MAX_ID)

    # let a client "route" to the node in question (accepts indices, which are analogous to routing information)
    def access_node(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID]
        else:
            return None

    def get_num_rounds_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_num_rounds()
        else:
            return None

    def get_recent_num_rounds_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_recent_num_rounds()
        else:
            return None

    def get_num_rounds(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_num_rounds(nodeID)
        else:
            return None

    def get_recent_num_rounds(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_recent_num_rounds(nodeID)
        else:
            return None

    def get_num_messages_sent_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_num_messages_sent()
        else:
            return None

    def get_recent_num_messages_sent_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_recent_num_messages_sent()
        else:
            return None

    def get_num_messages_sent(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_num_messages_sent(nodeID)
        else:
            return None

    def get_recent_num_messages_sent(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_recent_num_messages_sent(nodeID)
        else:
            return None

    def get_num_messages_recv_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_num_messages_recv()
        else:
            return None

    def get_recent_num_messages_recv_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_recent_num_messages_recv()
        else:
            return None

    def get_num_messages_recv(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_num_messages_recv(nodeID)
        else:
            return None

    def get_recent_num_messages_recv(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_recent_num_messages_recv(nodeID)
        else:
            return None

    def get_num_bytes_sent_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_num_bytes_sent()
        else:
            return None

    def get_recent_num_bytes_sent_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_recent_num_bytes_sent()
        else:
            return None

    def get_num_bytes_sent(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_num_bytes_sent(nodeID)
        else:
            return None

    def get_recent_num_bytes_sent(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_recent_num_bytes_sent(nodeID)
        else:
            return None

    def get_num_bytes_recv_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_num_bytes_recv()
        else:
            return None

    def get_recent_num_bytes_recv_base(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].get_recent_num_bytes_recv()
        else:
            return None

    def get_num_bytes_recv(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_num_bytes_recv(nodeID)
        else:
            return None

    def get_recent_num_bytes_recv(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_recent_num_bytes_recv(nodeID)
        else:
            return None

    # ONLY QP Nodes and higher will have these

    def get_finger_table_range_accesses(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_finger_table_range_accesses(nodeID)
        else:
            return None

    def get_finger_table_accesses(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_finger_table_accesses(nodeID)
        else:
            return None

    # ONLY QPLastHop Nodes will have these

    def get_database_accesses(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_database_accesses(nodeID)
        else:
            return None

    # ONLY DHTPIR Nodes will have these

    def get_PHF_generations(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_PHF_generations(nodeID)
        else:
            return None

    def get_PIR_retrievals(self, quorumID, nodeID):
        if quorumID >= 0 and quorumID < len(self.nodes):
            return self.nodes[quorumID].get_PIR_retrievals(nodeID)
        else:
            return None

    # End insert for calcuations

    def __access_node_tables(self, nodeID):
        if nodeID >= 0 and nodeID < len(self.nodes):
            return self.nodes[nodeID].table
        else:
            return None

    def get_num_nodes(self):
        return len(self.nodes)

    def test_tables(self):
        if self.get_num_nodes() == 10:
            retval = True
            table = self.__access_node_tables(5)
            retval = retval and table[b'\x99\x99\x99\x99\x99\x99\x98\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 6
            retval = retval and table[b'\xb3333330\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 7
            retval = retval and table[b'\xe6fffffh\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 9
            retval = retval and table[b'L\xcc\xcc\xcc\xcc\xcc\xcc\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 3
            table = self.__access_node_tables(9)
            retval = retval and table[b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 0
            retval = retval and table[b'\xb3333330\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 7
            retval = retval and table[b'\x19\x99\x99\x99\x99\x99\x9a\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 1
            retval = retval and table[b'L\xcc\xcc\xcc\xcc\xcc\xcc\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'] == 3
            return retval
        else:
            print("Assumed testing conditions not met; no tests run.")
            return False

# Normally this file is a class to be used elsewhere,
# but if you run it directly it performs some rudimentary unit tests
if __name__ == "__main__":
    from base_node import Base_Node

    SIZE_OF_DOCUMENTS_IN_TEST = 1024
    NUM_NODES_IN_TEST = 10
    
    # The 1 here isn't used at all within the simulator,
    # but it is customary to invoke DHT_Simulator with that value there for Base_Node
    test = DHT_Simulator(Base_Node, NUM_NODES_IN_TEST, SIZE_OF_DOCUMENTS_IN_TEST, 1)

    assert test.index_to_owner_node(8) == b'\xcc\xcc\xcc\xcc\xcc\xcc\xd0\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
    assert test.index_to_owner_node(4) == b'ffffffh\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
    assert test.owner_node_to_index(b'3333334\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00') == 2
    assert test.owner_node_to_index(b'\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00') == 5
    print("Index to nodeID conversion functioning correctly.")

    assert test.test_tables()
    print("Finger tables correctly generated.")