#
# Copyright (c) Daniel Sheffield 2023
# All rights reserved
#
# THIS SOFTWARE IS PROVIDED AS IS WITHOUT WARRANTY
from itertools import combinations
import random
import time
from urllib.parse import urlencode
from pytest import fixture, mark
import app.rest.hash_util as hash_util
import base32_lib as b32
from app.rest.hash_util import(
    hash_to_hex,
    hash_to_bytes,
    bytes_to_base32,
    bytes_to_base64,
    bytes_to_base85,
    bytes_to_hash,
    bytes_to_hex,
    normalize_bytes,
    normalize_hash,
    normalize_hex,
    base85_to_hash,
    base85_to_bytes,
    base85_to_hex,
    base32_to_bytes,
    base32_to_hash,
    base32_to_hex,
    base64_to_bytes,
    base64_to_hash,
    base64_to_hex,
    hash_to_base32,
    hash_to_base85,
    hash_to_base64,
    hex_to_base32,
    hex_to_base85,
    hex_to_base64,
    hex_to_bytes,
    hex_to_hash,
    python,
    shake,
    blake,
)

def to_unsigned(_hash: int) -> int:
    return _hash & hash_util.DIGEST_SIZE_BITMASK

def digest_size():
    for g in ( two_bytes, three_bytes, four_bytes):
        yield g

def two_bytes():
    hash_util.DIGEST_SIZE_BYTES = 2
    hash_util.DIGEST_SIZE_NIBBLES = 4
    hash_util.DIGEST_SIZE_BITMASK = 0xffff
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BIT = 0x10000
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BITMASK = 0x1ffff
    return hash_util.DIGEST_SIZE_BYTES

def three_bytes():
    hash_util.DIGEST_SIZE_BYTES = 3
    hash_util.DIGEST_SIZE_NIBBLES = 6
    hash_util.DIGEST_SIZE_BITMASK = 0xffffff
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BIT = 0x1000000
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BITMASK = 0x1ffffff
    return hash_util.DIGEST_SIZE_BYTES

def four_bytes():
    hash_util.DIGEST_SIZE_BYTES = 4
    hash_util.DIGEST_SIZE_NIBBLES = 8
    hash_util.DIGEST_SIZE_BITMASK = 0xffffffff
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BIT = 0x100000000
    hash_util.DIGEST_SIZE_SIGNED_TO_UNSIGNED_BITMASK = 0x1ffffffff
    return hash_util.DIGEST_SIZE_BYTES

@fixture
def two_byte_bitmask():
    return 0xffff

def hash_codes():
    for _hash in [
        -7583489610679606711,
        1421958803217889556,
        10,
    ]:
        _bytes = b'\00' * 8 + _hash.to_bytes(8, byteorder='big', signed=True)[-8:]
        _hex = _bytes.hex()
        print(_hash, _hex, _bytes)
        yield _hash, _hex, _bytes

# PYTHONHASHSEED=0

def digest_size_bitmask(digest_size_bits):
    mask = 0x0
    for _ in range(digest_size_bits):
        mask = mask << 1 | 0x1
    return mask

def digest_size_force_unsigned_bitmask(digest_size):
    mask = 0x1
    for _ in range(8*digest_size):
        mask = mask << 1 | 0x1
    return mask

def test_three_byte_assertions():
    digest_size = three_bytes()

    assert hash_to_base64(-7583489610679606711) == 'k6JJ'
    assert bytes_to_base64(b'\x96\xc2\x08`\xcd\x93\xa2I') == 'k6JJ'
    assert hex_to_base64('96c20860cd93a249') == 'k6JJ'
    assert base64_to_hex('k6JJ') == '93a249'
    assert base64_to_bytes('aa') == b'\x00\x00i'
    assert base64_to_hex('aa') == '000069'
    assert base64_to_hash('aa') == 105


@mark.parametrize('_hash, _hex', [
    (_hash, _hex) for _hash, _hex, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_hash_to_hex(_hash, _hex, digest_size_gen):
    digest_size = digest_size_gen()
    assert hash_to_hex(_hash) == _hex[-digest_size*2:]
    for i in range(1, digest_size*2+1):
        # test hex is zero padded
        assert hash_to_hex(_hash & digest_size_bitmask(i*4)) == _hex[-i:].zfill(digest_size*2)

@mark.parametrize('_hash, _bytes', [
    (_hash, _bytes) for _hash, _, _bytes, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_hash_to_base32(_hash, _bytes, digest_size_gen):
    digest_size = digest_size_gen()
    norm = hash_to_base32(_hash)
    assert norm == b32.encode(
        to_unsigned(int.from_bytes(
            _bytes[-digest_size:], byteorder='big', signed=False)
        )
    ).upper().zfill(digest_size*8//5+1)
    assert len(norm) == digest_size*8//5 +1
    for i in range(1, digest_size+1):
        norm = hash_to_base32(_hash & digest_size_bitmask(i*8))
        # test base32 is zero padded
        assert norm == b32.encode(
            to_unsigned(int.from_bytes(
                (b'\x00'*digest_size + _bytes[-i:])[-digest_size:], byteorder='big', signed=False)
            )
        ).upper().zfill(digest_size*8//5+1)
        assert len(norm) == digest_size*8//5 +1

@mark.parametrize('_hash, _hex', [
    (_hash, _hex) for _hash, _hex, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_hex_to_hash(_hash, _hex, digest_size_gen):
    digest_size = digest_size_gen()
    assert hex_to_hash(_hex) == _hash & digest_size_bitmask(digest_size*8)
    for i in range(1, digest_size*2+1):
        # test non-zero padded hex
        _input = _hex[-i:]
        assert digest_size*2 == i or len(_input) < digest_size*2
        assert hex_to_hash(_input) == _hash & digest_size_bitmask(i*4)

        # test zero padded hex
        _input = _hex[-i:].zfill(digest_size*2)
        assert len(_input) == digest_size*2
        assert hex_to_hash(_input) == _hash & digest_size_bitmask(i*4)

@mark.parametrize('_hex, _bytes', [
    (_hex, _bytes) for _, _hex, _bytes, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_hex_to_bytes(_hex, _bytes, digest_size_gen):
    digest_size = digest_size_gen()
    assert hex_to_bytes(_hex) == _bytes[-digest_size:]
    for i in range(1, digest_size*2+1):
        # test non-zero padded hex
        _input = _hex[-i:]
        assert digest_size*2 == i or len(_input) < digest_size*2
        if i % 2:
            if i > 2:
                assert hex_to_bytes(_input) == (b'\x00'*digest_size + bytes([
                    _bytes[-i//2] & 0x0f, *_bytes[-i//2+1:]
                ]))[-digest_size:]
            else:
                assert hex_to_bytes(_input) == (b'\x00'*digest_size + bytes([
                    _bytes[-i//2] & 0x0f
                ]))[-digest_size:]
        else:
            assert hex_to_bytes(_input) == (b'\x00'*digest_size + _bytes[-i//2:])[-digest_size:]

        # test zero padded hex
        _input = _hex[-i:].zfill(digest_size*2)
        assert len(_input) == digest_size*2
        if i % 2:
            if i > 2:
                assert hex_to_bytes(_input) == (b'\x00'*digest_size + bytes([
                    _bytes[-i//2] & 0x0f, *_bytes[-i//2+1:]
                ]))[-digest_size:]
            else:
                assert hex_to_bytes(_input) == (b'\x00'*digest_size + bytes([
                    _bytes[-i//2] & 0x0f
                ]))[-digest_size:]
        else:
            assert hex_to_bytes(_input) == (b'\x00'*digest_size + _bytes[-i//2:])[-digest_size:]

@mark.parametrize('_hash, _bytes', [
    (_hash, _bytes) for _hash, _, _bytes, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_bytes_to_hash(_hash, _bytes, digest_size_gen):
    digest_size = digest_size_gen()
    assert bytes_to_hash(_bytes) == _hash & digest_size_bitmask(digest_size*8)
    for i in range(1, digest_size+1):
        # test non-zero padded bytes
        _input = _bytes[-i:]
        assert digest_size == i or len(_input) < digest_size
        assert bytes_to_hash(_input) == _hash & digest_size_bitmask(i*8)

        # test zero padded bytes
        _input = (b'\x00'*digest_size + _bytes[-i:])[-digest_size:]
        assert len(_input) == digest_size
        assert bytes_to_hash(_input) == _hash & digest_size_bitmask(i*8)

@mark.parametrize('_hash, _bytes', [
    (_hash, _bytes) for _hash, _, _bytes, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_base32_to_hash(_hash, _bytes, digest_size_gen):
    digest_size = digest_size_gen()
    _base32 = b32.encode(
        to_unsigned(int.from_bytes(
            _bytes[-digest_size:], byteorder='big', signed=False)
        )
    ).upper().zfill(digest_size*8//5+1)
    norm = base32_to_hash(_base32)
    assert norm == _hash & digest_size_bitmask(digest_size*8)

    for i in range(1, digest_size+1):
        # test non-zero padded base32
        _base32 = b32.encode(
            to_unsigned(int.from_bytes(
                (b'\x00'*digest_size + _bytes[-i:])[-digest_size:], byteorder='big', signed=False)
            )
        ).upper()
        norm = base32_to_hash(_base32)
        assert norm == _hash & digest_size_bitmask(i*8)

        # test zero padded base32
        _base32 = b32.encode(
            to_unsigned(int.from_bytes(
                (b'\x00'*digest_size + _bytes[-i:])[-digest_size:], byteorder='big', signed=False)
            )
        ).upper().zfill(digest_size*8//5+1)
        norm = base32_to_hash(_base32)
        assert norm == _hash & digest_size_bitmask(i*8)

@mark.parametrize('_hex', [
    _hex for _, _hex, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_normalize_hex(_hex, digest_size_gen):
    digest_size = digest_size_gen()
    assert normalize_hex(_hex) == _hex[-digest_size*2:]
    assert len(normalize_hex(_hex)) == digest_size*2

    # test commutivity
    for i in range(1, digest_size*2+1):
        assert i == digest_size*2 or _hex[-i:] != _hex[-i:].zfill(digest_size*2)
        assert normalize_hex(_hex[-i:]) == _hex[-i:].zfill(digest_size*2)
        assert len(normalize_hex(_hex[-i:])) == digest_size*2

@mark.parametrize('_bytes', [
    _bytes for _, _, _bytes, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_normalize_bytes(_bytes, digest_size_gen):
    digest_size = digest_size_gen()
    assert normalize_bytes(_bytes) == _bytes[-digest_size:]
    assert len(normalize_bytes(_bytes)) == digest_size
    
    # test commutivity
    for i in range(1, digest_size+1):
        assert i == digest_size or _bytes[-i:] != (b'\x00'*digest_size + _bytes[-i:])[-digest_size:]
        assert normalize_bytes(_bytes[-i:]) == (b'\x00'*digest_size + _bytes[-i:])[-digest_size:]
        assert len(normalize_bytes(_bytes[-i:])) == digest_size

@mark.parametrize('_hash', [
    _hash for _hash, *_ in hash_codes()
])
@mark.parametrize('digest_size_gen', digest_size())
def test_normalize_hash(_hash, digest_size_gen):
    digest_size = digest_size_gen()
    norm = normalize_hash(_hash)
    assert norm == _hash & digest_size_bitmask(digest_size*8)
    
    # test commutivity
    for i in range(1, digest_size+1):
        norm = normalize_hash(_hash & digest_size_bitmask(i*8))
        assert  norm == _hash & digest_size_bitmask(i*8)
        assert norm == norm & digest_size_bitmask(i*8)

#
# Performance tests

four_bytes()
products = { k for k in """
Adzuki Beans|Alberts Eggs|Almonds|Apple Cider Vinegar|Apples|Apricots|Arrowroot Powder|Avocado|Baking Soda|Bananas|Barn Eggs|Bars|Beans Canned|Beef Chuck|Beef Jerky|Beef Mince|Beef Patties|Beef Roast|Beef Sausages|Beef Scotch Fillet|Beef Shin|Beef Stir Fry|Beef Topside|Beetroot|Beetroot Canned|Blade Steak|Blueberries Frozen|Brazil Nuts|Brewers Yeast|Brisket|Broccoli|Brown Lentils|Brown Rice|Brown Sugar|Buckwheat Groats|Buns|Burger Patties|Butter|Butternut Squash|Cabbage|Cacao Powder|Cage Free Eggs|Candied Ginger|Canned Apricots|Canned Peaches|Canned Salmon|Canned Tomatoes|Canned Tuna|Capsicum|Cardamon|Carrots|Cashews|Cauliflower|Cayenne Powder|Celery|Chicken Drumsticks|Chicken Frames|Chicken Mince|Chicken Nibbles|Chicken Stir Fry|Chicken Tender|Chicken Thighs|Chicken Thighs BL|Chickpeas|Chocolate Chips|Cinnamon, Ground|Coconut Cream|Coconut Desiccated|Coconut Oil|Coconut Water|Coconut Yoghurt|Coffee Beans|Colby Cheese|Cookie Pack|Cookie Time|Corn Chips|Corned Beef|Cornflour|Courgettes|Crackers|Cream|Cruskits|Cucumber|Curry Paste|Dark Chocolate|Demerara Sugar|Doritos|Dried Dates|Dried Figs|Dry Peas|Edam Cheese|Eggplant|FR Chicken Breast|FR Chicken Drumsticks|FR Chicken Mince|FR Chicken Thighs|FR Chicken Whole|Feta|Fish Fillets|Fish Frames|Fish Heads|Fish Sauce|Flax Seeds|Free Range Eggs|French Lentils|Frozen Cherries|Frozen Green Beans|Frozen Peas|Frozen Pop|Fruit Bars|Garam Masala|Garlic|Ginger|Ginger Beer|Ginger Ground|Hemp Seeds|Herbal Tea|Honey|Hot Chilli|Hummus|Ice Cream|Jalapenos Canned|Jar of Chillis|Kassori Methi|Kelp|Kiwifruit|Kumara|Lamb Cubes|Lamb Mince|Lamb Pieces|Lamb Shank|Lamb Shoulder|M&Ms Mini|Magnesium Fizz|Mandarins|Mango|Maple Syrup|Marmite|Mayo|Meatballs|Mild Cheese|Milk Standard|Mixed Fruit Frozen|Mixed Nuts|Molasses|Monkfruit|Mozzarella Cheese|Mung Beans|Mushrooms|Mutton Chops|Neck Chops|Nut Bars|Nutritional Yeast|OSM bars|Oat Bars|Olive Oil|Olives Bag|Olives Jar|Onions|Orange Juice|Oranges|Oreo Cookies|Organic Cheese|Organic Chicken|Paprika|Parsnips|Pasta|Peaches|Peanutbutter|Peanuts|Pears|Pinenuts|Pizza Bases|Plums|Potato Crisps|Potatoes|Powdered Milk|Prunes|Puffed Rice|Pukka Tea|Pumpkin Seeds|Quinoa|Raw Milk|Red Lentils|Rice Crackers|Rice Noodles|Rooibos Tea|Rye Flour|Ryecorn|Salad Bag|Salami|Salsa|Salt|Saurkraut|Seaweed|Smoked Salmon|Snack Balls|Soda Syrup|Soup Refrigerated|Soy Sauce|Squash|Starter|Strawberries|Strawberries Frozen|Sugar, Coconut|Sultanas|Sunflower Oil|Sunflower Seeds|TNCC Lollies|Tasty Cheese|Tomato Fresh|Tomato Juice|Tomato Paste|Tomato Sauce|Tonic Water|Tulsi Tea|Tumeric|Vanilla|Veggie Crisps|Vinegar|Vintage Cheese|Walnut|Water|Watermelon|Wheat Berries|White Flour|White Rice|White Sugar|Whittakers Chocolate|Whole Chicken|Whole Fish|Whole Oats|Wholemeal Flour|Xcut Blade Steak|Yoghurt Dairy
""".strip().split('|')
}

def find_collision(f=python, norm=normalize_hash, how_many=None, **kwargs):
    seen = dict()
    random.shuffle(how_many)
    for i in how_many:
        for idx, j in enumerate(combinations(products, i)):
            if idx > 10000:
                break
            url = urlencode([
                ('product', k) for k in j
            ])
            _hash = f(url.encode('utf-8'), **kwargs).hex()
            _hash = norm(_hash) if callable(norm) else _hash
            if _hash in seen:
                return url, _hash, seen
            #assert _hash not in seen, len(seen)
            #assert _hash not in seen, f"{url} ({_hash}) == {seen[_hash]} ({len(seen)})"
            seen[_hash] = url
    return None, None, seen

# trials = 5 #20
# for name, args, kwargs in [
#     #("Python", [], {'norm': normalize_hash}),
#     #("Shake", [shake], {'norm': hex_to_hash}),
#     ("Blake", [blake], {'norm': hex_to_hash, 'person': 'trend'.encode("utf-8")}),
#     #("Python", [], {'norm': hash_to_base32}),
#     #("Shake", [shake], {'norm': hex_to_base32}),
#     #("Blake", [blake], {'norm': hex_to_base32, 'person': 'trend'.encode("utf-8")}),
# ]:
#     start = time.time()
#     random.seed(0)
#     score = 0
#     _min = 0
#     print(name)
#     for i in range(trials):
#         url, _hash, seen = find_collision(*args, **kwargs, how_many=list(range(len(products))))
#         n = len(seen)
#         print(n)
#         if n < 30000:
#             print(url)
#             print(seen[_hash])
#         score += n
#         _min = _min if _min < n and _min else n
#     end = time.time()
#     print(f"Score: {_min} / {score/trials}")
#     print(f"Time: {end - start}")

# exit(0)


if __name__ == '__main__':
    import timeit
    # For Python>=3.5 one can also write:
    #print(timeit.timeit("normalize_bytes(b'a')", globals=locals(), number=5000000))
    for method, arg in [
        ('normalize_hash', "hash('a')"),
        ('normalize_hex', "'a'"),
        ('normalize_bytes', r"b'\x0a'"),
        #('hex_to_hash', "'a'"),
        #('hex_to_bytes', "'a'"),
        #('bytes_to_hash', r"b'\n'"),
        #('bytes_to_hex', r"b'\n'"),
        #('hash_to_bytes', "10"),
        #('hash_to_hex', "10"),
        # ('base64_to_hash', "'aa'"),
        # ('base64_to_bytes', "'aa'"),
        # ('base64_to_hex', "'aa'"),
        # ('bytes_to_base64', r"b'\n'"),
        # ('hex_to_base64', "'a'"),
        # ('hash_to_base64', "10"),
        # ('base85_to_hash', "'aa'"),
        # ('base85_to_bytes', "'aa'"),
        # ('base85_to_hex', "'aa'"),
        # ('bytes_to_base85', r"b'\n'"),
        # ('hex_to_base85', "'a'"),
        # ('hash_to_base85', "10"),
        ('base32_to_hash', "'aa'"),
        ('base32_to_bytes', "'aa'"),
        ('base32_to_hex', "'aa'"),
        ('bytes_to_base32', r"b'\n'"),
        ('hex_to_base32', "'a'"),
        ('hash_to_base32', "10"),
        #('shake', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
        #('blake', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
        #('md5hash', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
        #('sha256hash', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
        #('sha1hash', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
        #('python', "'abc1234567890&|abcdefghijklmnopqrstuvwxyz'"),
    ]:
        test = f"{method}({arg})"
        elapsed = timeit.timeit(test, globals=locals(), number=5000000)
        print(f"{method}: {elapsed}")