import hashlib
from binascii import hexlify,unhexlify
from os import urandom
from math import ceil,log,floor
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 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

    
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 random_key(n=32):    
    return hexlify(urandom(n))


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

def fn_k(x,k):
    return hash(k.decode()+x)


def chain_fn(x,r,i,k):
    if i == 0:
        return x
    else:
        for y in range(i):
            x = fn_k(hex(int(x,16)^int(r[y],16))[2:-1],k)
    return x


def chain_fn2(x,r,i,k):
        for y in range(i,15):
            x = fn_k(hex(int(x,16)^int(r[y],16))[2:-1],k)
        return x


def random_wpkey(w=16):
    # first calculate l_1 + l_2 = l
    # if using SHA-256 then m and n = 256

    if w==16:
        l=67
        l_1=64
        l_2=3
    else:
        m = 256
        l_1 = ceil(m/log(w,2))
        l_2 = floor(log((l_1*(w-1)),2)/log(w,2)) + 1
        l = int(l_1+l_2)

    sk = []
    pub = []

    # create l+w-1 256 bit secret key fragments
    # l n-bits will be private key, remaining w-1 will be r, the randomisation elements for the chaining function
    # finally generate k the key for the chaining function

    for x in range(l+w-1):
        sk.append(random_key())

    priv = sk[:-(w-1)]
    r = sk[l:]

    k = random_key()

    pub.append((r,k))       # pk_0 = (r,k) where r is a list of w-1 randomisation elements

    for sk_ in priv:
        pub.append(chain_fn(sk_,r,w-1,k))

    return priv, pub


def sign_wpkey(priv, message, pub, w=16):            

    m = 256
    if w==16:
        l_1 = 64
        l_2 = 3
        l = 67
    else:    
        l_1 = ceil(m/log(w,2))
        l_2 = floor(log((l_1*(w-1)),2)/log(w,2)) + 1
        l = int(l_1+l_2)

    message = hash(message)       # outputs 256 bit -> 64 hex, each hex digit is therefore 4 bits

    s = []
    checksum = 0

    for x in range(int(l_1)):
        y = int(message[x],16)
        s.append(y)
        checksum += w-1-y

    c = (hex(checksum))[2:]

    if len(c) < 3:
        c = '0'+c

    for x in range(int(l_2)):
        y = int(c[x],16)
        s.append(y)

    signature = []

    for x in range(int(l)):        
        signature.append(chain_fn(priv[x],pub[0][0],s[x],pub[0][1]))

    return signature


def verify_wpkey(signature, message, pub, w=16):
    m = 256
    if w==16:
        l_1 = 64
        l_2 = 3
        l = 67
    else:    
        l_1 = ceil(m/log(w,2))
        l_2 = floor(log((l_1*(w-1)),2)/log(w,2)) + 1
        l = int(l_1+l_2)

    message = hash(message)

    s = []
    checksum = 0

    for x in range(int(l_1)):
        y = int(message[x],16)
        s.append(y)
        checksum += w-1-y

    c = (hex(checksum))[2:]

    if len(c) < 3:
        c = '0'+c

    for x in range(int(l_2)):
        y = int(c[x],16)
        s.append(y)

    pub2 = []

    for x in range(int(l)):         
        pub2.append(chain_fn2(signature[x],pub[0][0],s[x],pub[0][1]))

    if pub2 == pub[1:]:
        return True

    return False


