#! /usr/bin/env python

import io, random, sys

import importlib.machinery, importlib.util

def importhack():
    global base85
    loader = importlib.machinery.SourceFileLoader("base85", "./base85")
    spec = importlib.util.spec_from_loader("base85", loader)
    base85 = importlib.util.module_from_spec(spec)
    loader.exec_module(base85)

def main():
    importhack()
    random.seed(0x29F8A6D4)
    testdata = random.randbytes(262144)
    testrange = []

    testrange.extend([i for i in range(64)])

    testrange.append(base85.MAX_DECODE_BUFFER_LEN - 1)
    testrange.append(base85.MAX_DECODE_BUFFER_LEN)
    testrange.append(base85.MAX_DECODE_BUFFER_LEN + 1)
    testrange.append(base85.MAX_ENCODE_BUFFER_LEN - 1)
    testrange.append(base85.MAX_ENCODE_BUFFER_LEN)
    testrange.append(base85.MAX_ENCODE_BUFFER_LEN + 1)

    testrange.append(2 * base85.MAX_DECODE_BUFFER_LEN - 1)
    testrange.append(2 * base85.MAX_DECODE_BUFFER_LEN)
    testrange.append(2 * base85.MAX_DECODE_BUFFER_LEN + 1)
    testrange.append(2 * base85.MAX_ENCODE_BUFFER_LEN - 1)
    testrange.append(2 * base85.MAX_ENCODE_BUFFER_LEN)
    testrange.append(2 * base85.MAX_ENCODE_BUFFER_LEN + 1)

    testrange.append(3 * base85.MAX_DECODE_BUFFER_LEN - 1)
    testrange.append(3 * base85.MAX_DECODE_BUFFER_LEN)
    testrange.append(3 * base85.MAX_DECODE_BUFFER_LEN + 1)
    testrange.append(3 * base85.MAX_ENCODE_BUFFER_LEN - 1)
    testrange.append(3 * base85.MAX_ENCODE_BUFFER_LEN)
    testrange.append(3 * base85.MAX_ENCODE_BUFFER_LEN + 1)

    datalength, rest = divmod(base85.MAX_ENCODE_BUFFER_LEN, 5)
    assert rest == 0
    testrange.append(datalength * 8 - 1)
    testrange.append(datalength * 8)
    testrange.append(datalength * 8 + 1)

    datalength, rest = divmod(base85.MAX_DECODE_BUFFER_LEN, 5)
    assert rest == 0
    testrange.append(datalength * 8 - 1)
    testrange.append(datalength * 8)
    testrange.append(datalength * 8 + 1)

    testrange.append(10_000)
    testrange.append(100_000)
    testrange.append(131072)
    testrange.append(262144)
    for i in range(200):
        testrange.append(random.randint(0, 262144))

    for n in testrange:
        assert 0 <= n <= 262144

    for n in testrange:
        src = io.BytesIO(testdata[:n])
        dst = io.BytesIO()
        base85.encode(src, dst)
        ret_src = io.BytesIO(dst.getvalue())
        ret_dst = io.BytesIO()
        base85.decode(ret_src, ret_dst, False)
        if ret_dst.getvalue() != testdata[:n]:
            print("Failed at length =", n)
            sys.exit(1)

if __name__ == "__main__":
    main()