import hashlib
from binascii import hexlify,unhexlify
from os import urandom
import sha3


def dimension_array(N,B):
    array = []
    for i in range(N):
        array_N = []
        for j in range(B):
            array_N.append("")
        array.append(array_N)
    return array    


def one_dimension_array(N):
    array = []
    for i in range(N):
        array_N = []
        array.append(array_N)
    return array  

    
def hash(message, hash_digest: int = 256):
    if hash_digest == 256:
        return hashlib.sha3_256(str(message).encode('utf-8')).hexdigest()
    else:
        # possibility to implement bigger hashes
        raise ValueError('Different digest size hashes not implemented as of now.')


def leaf_to_string(leaf):

    index, nonce = leaf[1]
    public_key_hash = hash(leaf[0])

    helper_string = public_key_hash + ";" + index.decode() + ";" + nonce.decode()
    helper_string = helper_string.encode().hex() 

    return  helper_string


def string_to_leaf(leaf_hex_string):
    leaf_array = bytes.fromhex(leaf_hex_string).decode().split(";")
    leaf_array[1] = leaf_array[1].encode()
    leaf_array[2] = leaf_array[2].encode()

    return  [leaf_array[0],(leaf_array[1],leaf_array[2])]


def get_leaf_from_signature(signature):
    pk = signature[1]
    encrypted_index = signature[2]
    leaf = leaf_to_string([pk,encrypted_index])
    
    return leaf


################################################### OTS ###################################################
#                                                                                                         #
#           OTS keygen, sign, verify operations modified from https://github.com/surg0r/lamport           #
#                                                                                                         #
################################################## WOTS ###################################################

def random_key(n=32):                 
    return hexlify(urandom(n))


def random_wkey(w=8, verbose=0):      #create random W-OTS keypair
    if w > 16:
        w = 16      #too many hash computations to make this sensible
    priv = []
    pub = []

    for x in range(256//w):
        a = random_key()
        priv.append(a)
        for y in range(2**w-1):             
            a = hash(a)
        pub.append(hash(a))              

    return priv, pub    


def sign_wkey(priv, message):      

    signature = []
    bin_msg = unhexlify(hash(message))

    for y in range(len(priv)):
        s = priv[y]    
        for x in range(256-ord(bin_msg[y:y+1])):
            s = hash(s)
        signature.append(s)

    return signature
    

def verify_wkey(signature, message, pub):

    verify = []
    bin_msg = unhexlify(hash(message))
    
    for x in range(len(signature)):
        a = signature[x]
                                                   
        for z in range(ord(bin_msg[x:x+1])):
                a=hash(a)                          
        verify.append(a)
  
    if pub != verify:
        return False

    return True
