import entities
import utils
import sys
from math import floor
from time import time

class g_merkle_GSS:
    def __init__(self, 
                    N: int = 4,             
                    B: int = 4,             
                    #hash_digest: int = 256  # possible future extension to implement stronger hash functions
                    ):
        '''
        Initialize the G-Merkle Group Signature Scheme with given parameters.

        Args:
            N:      number of members
            B:      number of signatures per member    

        Also initializes group manager, joins members, builds Merkle tree and assigns group public key as Merkle root.
        '''
        t = time()
        self.N = N
        self.B = B
        self.num_of_leaves = B*N
        #self.hash_digest = hash_digest

        self.gm = entities.group_manager(self.N,self.B)
        self.gpk = self.gm.gpk
        print(f"\nScheme instantiation took {time()-t}s in total.\n")

    def sign(self, member_id: int, message: str):
        '''
        User with given id can sign a given message.

        Args:
            member_id:      member's id from range 1 - N
            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
            signature[3]:   authentication path
            signature[4]:   message hash           
        '''

        sys.tracebacklimit = 0
        t = time()
        print(f"Signing a message ({len(message)} B)...",end="", flush=True)

        # check member id and invoke member sign
        if -1 < member_id-1 < self.N:
            try:
                signature_and_auth_path = self.gm.members[member_id-1].sign(message)
            except ValueError:
                print(f" Member with id {member_id} cannot sign any more messages.")
                return False
        else:
            print (f' Member with id {member_id} is not a part of the group.')
            return

        # check if public key is present within the tree leaves and get its ID
        is_pk_present_in_the_tree_leaves = self.check_leaves_for_public_key_hash(signature_and_auth_path)

        if (isinstance(is_pk_present_in_the_tree_leaves,int)):
            signature_and_auth_path.append(self.get_authentication_path(int(is_pk_present_in_the_tree_leaves)))
            signature_and_auth_path.append(utils.hash(message))
        else:
           print(is_pk_present_in_the_tree_leaves)
           return

        print(f"\t\t\t\t\t\t[DONE] {time()-t}s")

        return signature_and_auth_path


    def verify(self, signature, message: str, gpk):
        '''
        Verification of given signature over the group public key.
        Consists of underlying OTS scheme verification and authentication path verification.

        Args:
            signature:      output from the sign() function
            message:        string message used to verify signature
            gpk:            group public key, Merkle tree root

        Returns:
            True            if both verifications are valid
            False           if either or both verifications are not valid
        '''
        t = time()
        print(f"Verifying signature...",end="", flush=True)

        try:
            #verify ots scheme
            underlying_ots_verification = self.ots_verify(signature,message)
            # verify authentication path and merkle root
            merkle_verification = self.validate_merkle_proof(signature[3],utils.get_leaf_from_signature(signature),gpk)
            print(f"\t\t\t\t\t\t\t[DONE] {time()-t}s")

            return underlying_ots_verification and merkle_verification
        except TypeError:
            print(" No signature for verification was given.")


    def open(self, signature, message: str):
        '''
        Opens a given signature.

        Args:
            signature:      output from the sign() function
            message:        string message to open signature

        Returns:
            string          output depending on the success of opening the signature
        '''
        t = time()
        print(f"Opening signature...",end="", flush=True)

        try:
            identificator = signature[2]
            message_hash = signature[4]

            if utils.hash(message) == message_hash:
                searched_leaf = utils.get_leaf_from_signature(signature)
                for leaf in self.gm.string_leaves:
                    if leaf[0] == searched_leaf:
                        id = self.gm.decrypt_index(identificator)
                        id = floor(float(id)/self.B)
                        print(f"\t\t\t\t\t\t\t[DONE] {time()-t}s\n")
                        print()
                        return f"Signature was created by user {id+1}."
                print(f"\t\t\t\t\t\t\t[DONE] {time()-t}s\n")
                return "Signature was not created by any valid group member."
            
            else:
                print(f"\t\t\t\t\t\t\t[DONE] {time()-t}s\n")
                return "Message hashes do not match."
        except TypeError:
            print(" No signature for opening was given.")
            print("\n")


###################################### HELPER FUNCTIONS ######################################

    def ots_verify(self, signature, message: str):
        sig = signature[0]
        public_key = signature[1]
        underlying_ots_verification = utils.verify_wpkey(sig, message, public_key)

        return underlying_ots_verification


    def get_authentication_path(self, id: int):

        return self.gm.merkle_tree.get_proof(id)


    def validate_merkle_proof(self, proof, hash, gpk):
        is_valid = self.gm.merkle_tree.validate_proof(proof,hash,gpk)

        return is_valid


    def check_leaves_for_public_key_hash(self, signature_and_auth_path):
        # check if given public key and leaf are present within the published tree
        leaf = utils.get_leaf_from_signature(signature_and_auth_path)

        for i in range(self.gm.merkle_tree.get_leaf_count()):
            if leaf == self.gm.merkle_tree.get_leaf(i):
                return i
    
        return 'Leaf was not found within the leaf nodes.'


