import merkletools
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import os, math, utils
import sha3
from base64 import b64decode,b64encode
from time import time


class group_member:

    def __init__(self, id: int, B: int):
        '''
        Initialize a group member.

        Args:
            id:     member's id
            B:      number of signatures one member can create

        Also generates B private/public keypairs for underlying OTS scheme.
        '''
        self.id = id
        self.B = B

        self.public_keys = []
        self.hashed_public_keys = []
        self.private_keys = []
        self.unencrypted_indexes = []
        self.encrypted_indexes = []
        self.leaves = []
        self.state = []     # states as given in the paper
        self.t = 0          # state counter

        self.generate_keypairs()


    def generate_keypairs(self):
        # generate B keypairs for self
        for _ in range (self.B):
            priv,pub = utils.random_wkey()
            self.public_keys.append(pub)
            self.private_keys.append(priv)


    def set_encrypted_indexes(self,some_array):
        self.encrypted_indexes = some_array


    def set_unencrypted_indexes(self,some_array):
        self.unencrypted_indexes = some_array


    def generate_leaves(self):
        for i in range(len(self.unencrypted_indexes)):
            # append public key and corresponding encrypted index and nonce depending on the order of public keys in members array
            self.leaves.append([utils.hash(self.public_keys[int(self.unencrypted_indexes[i])%self.B]),self.encrypted_indexes[i]])


    def generate_states(self):
        # create tuples matching ids and encrypted ids
        for i in range(len(self.unencrypted_indexes)):
            self.state.append((self.unencrypted_indexes[i],self.encrypted_indexes[i]))


    def sign(self, message: str):
        '''
        Invoke underlying OTS scheme's signature with public key determined by member's state.

        Args:
            message:        string message to be signed
        
        Returns:
            signature[0]:   signature on a given message
            signature[1]:   public key for a given signature
            signature[2]:   encrypted identificator for given public key
        
        Throws a value error if member cannot sign any more messages.

        '''
        # check if counter allows any more signature to be made
        if self.t < self.B:
            which_key, identificator = self.state[self.t]
            pk = self.public_keys[int(which_key)%self.B]
            sk = self.private_keys[int(which_key)%self.B]

            # invoke OTS signing process
            signature = utils.sign_wkey(sk,message)
            self.t += 1

            return  [signature, pk, identificator]
        else:
            raise ValueError('User cannot sign any more messages.')


class group_manager:

    def __init__(self, N: int, B: int):
        self.gpk = None
        self.merkle_tree = None
        gsk = os.urandom(32)
        self.gsk = gsk
        self.all_indexes = []
        self.leaves = []
        self.sorted_leaves = []
        self.string_leaves = []
        self.members = []
        self.N = N
        self.B = B
        self.num_of_leaves = N*B

        self.setup()


    def setup(self):
        # actions required for the scheme instantiation
        print(f"Joining {self.N} members and generating {self.num_of_leaves} OTS keypairs...",end="", flush=True)
        t = time()
        self.join_members()
        print(f"\t\t\t[DONE] {time()-t}s")
 
        print(f"Generating and shuffling {self.num_of_leaves} leaves and encrypting indexes...",end="", flush=True)
        t = time()
        self.encrypted_indexes_in_order()
        self.shuffled = self.shuffle_encrypted_indexes(self.all_indexes)
        self.assign_indexes()
        self.generate_leaves_states_for_members()
        self.generate_leaf_array()
        self.sort_leaf_array()
        self.all_leaves_to_string()
        print(f"\t\t[DONE] {time()-t}s")
       
        print(f"Building a Merkle tree with {self.num_of_leaves} leaves...",end="", flush=True)
        t = time()
        self.build_tree()
        print(f"\t\t\t\t[DONE] {time()-t}s")


    def join_members(self):
        # add members to the group
        for i in range (self.N):
            member = group_member(i,self.B)
            self.members.append(member)


    def encrypt_index(self,index):
        # encrypt plaintext identificator
        cipher_configuration = AES.new(self.gsk,AES.MODE_GCM)
        ciphertext = cipher_configuration.encrypt(pad(index.encode(),32))

        return b64encode(ciphertext),b64encode(cipher_configuration.nonce)


    def decrypt_index(self,encrypted_index):
        # decrypt encrypted identificator
        ciphertext, nonce = encrypted_index
        cipher = AES.new(self.gsk, AES.MODE_GCM, nonce = b64decode(nonce))

        return unpad(cipher.decrypt(b64decode(ciphertext)),32).decode()


    def encrypted_indexes_in_order(self):
        # array of arrays where first index corresponds to member id
        all_indexes_in_order = []
        counter = 0

        for _ in range (self.N):
            one_member_indexes = [] 
            for _ in range(self.B):
                index = self.encrypt_index(str(counter))
                one_member_indexes.append(index)
                # all indexes in one array
                self.all_indexes.append(index)
                counter += 1
            all_indexes_in_order.append(one_member_indexes)

        return all_indexes_in_order


    def shuffle_encrypted_indexes(self,encrypted_indexes):
        sorted_by_encrypted_indexes = sorted(encrypted_indexes, key=lambda tup: tup[0])

        return sorted_by_encrypted_indexes

    
    def assign_indexes(self):
        # assign indexes to members
        myarray = utils.one_dimension_array(self.N)
        myarray2 = utils.one_dimension_array(self.N)

        for tuple in self.shuffled:
            dec_member = self.decrypt_index(tuple)
            y_pos = math.floor(float(dec_member)/self.B)
            myarray[y_pos].append(tuple)
            myarray2[y_pos].append(dec_member)

        for i,member in enumerate(self.members):
            # assign tuple of encrypted index and nonce to member
            member.set_encrypted_indexes(myarray[i])
            # assign unencrypted indexes to members so that they can generate state
            member.set_unencrypted_indexes(myarray2[i])


    def generate_leaves_states_for_members(self):
        # tell members to generate leaves and states from encrypted indexes
        for member in self.members:
            member.generate_leaves()
            member.generate_states()


    def generate_leaf_array(self):
        # get leaves from members and generate leaf array
        helper_array = []

        for i in range(len(self.members)):
            for leaf in self.members[i].leaves:
                helper_array.append(leaf)

        self.leaves = helper_array


    def sort_leaf_array(self):
        # sort all leaves depending on encrypted identificator
        sorted_l = sorted(self.leaves, key=lambda tup: tup[1][0])
        self.sorted_leaves = sorted_l

    
    def all_leaves_to_string(self):
        # convert all leaves to hex strings
        helper_string = ""

        for i in range(len(self.sorted_leaves)):
            helper_string = self.sorted_leaves[i][0] + ";" + self.sorted_leaves[i][1][0].decode() + ";" + self.sorted_leaves[i][1][1].decode()
            helper_string = helper_string.encode().hex() 
            self.string_leaves.append([helper_string])
            

    def build_tree(self):
        # build Merkle tree
        self.merkle_tree = merkletools.MerkleTools(hash_type="sha3_256")

        for i in range(len(self.string_leaves)):
            self.merkle_tree.add_leaf(self.string_leaves[i][0])
        self.merkle_tree.make_tree()

        self.gpk = self.merkle_tree.get_merkle_root()

