import z3


def _ascii_printable(x):
    return z3.And(0x21 <= x, x <= 0x7e)  # enforces that byte is printable ascii


def _generate_ascii_printable_string(base_name, size, solver):
    # establishes z3 variable names for bytes of the input string
    bytes = [z3.BitVec('%s%d' % (base_name, i), 8) for i in range(size)]
    # adds the constraint that the bytes are printable ascii
    solver.add(z3.And([_ascii_printable(byte) for byte in bytes]))
    return bytes


def _str_to_BitVecVals8(string):  # transforms the target string into z3 BitVecVals
    return [z3.BitVecVal(ord(x), 8) for x in string]

def _get_collisions(hash_func, target, target_type, length, n_collisions, hash_table_size, *args):
    ret = []
    s = z3.Solver()
    # houses the z3 variables for the potential hash match
    res = _generate_ascii_printable_string('res', length, s)
    # enforces the z3 constraint that the hash matches the target
    # if the target_type is 'preimage', then we compare the hash to the hash of target
    if target_type == 'preimage':
        s.add(hash_func(res, hash_table_size, *args) == hash_func(_str_to_BitVecVals8(target), hash_table_size, *args))
        if length == len(target):  # don't generate the given preimage as an output if generating inputs of that length
            s.add(z3.Or([r != ord(t) for r, t in zip(res, target)]))
    # otherwise the target_type is 'image', and we compare the hash to the target itself
    else:
        s.add(hash_func(res, hash_table_size, *args) == target)
    count = 0
    # z3 isn't stateful; you have to run it again and again while adding constraints to ignore previous solutions
    while s.check() == z3.sat:
        x = s.model()  # This is a z3 solution
        y = ''.join(chr(x[i].as_long()) for i in res)
        ret.append(y)
        count += 1
        # add constraints
        s.add(z3.Or([r != x[r] for r in res]))
        if count >= n_collisions:
            ret.sort()
            break
    return ret

def get_collisions(hash_func, target, target_type, length, n_collisions, hash_table_size, *args):
    if length is None:
        ret = []
        i = 0
        while len(ret) < n_collisions:
            i = i + 1
            ret.extend(_get_collisions(hash_func, target, target_type, i,
                                       n_collisions - len(ret), hash_table_size, *args))
    else:
        ret = _get_collisions(hash_func, target, target_type, length, n_collisions, hash_table_size, *args)
        if len(ret) < n_collisions:
            # raises exception if the current number of solutions was not found
            raise ValueError('Could not find enough solutions')
    return ret
