from sampy.raknet.bitstream import Bitstream

bitstream = Bitstream(b"\x00" * 4)
MAGIC_BYTES = b"\x01\xff\x02"
MAGIC_BITS = [int(bit) for bit in " ".join([format(byte, "08b") for byte in MAGIC_BYTES]) if bit in ("0", "1")]
fail = 0

def print_pretty(prefix, msg):
    prefix += " " * (40 - len(prefix))
    print(prefix, msg)

def reset_bitstream():
    bitstream._buffer = bytearray(b"\x00" * len(bitstream._buffer))
    bitstream._offset = 0

def test_write_bit(var_bits):
    global fail
    print_pretty("<Test bitstream.write_bit:in>", bitstream.pretty())
    for bit in var_bits:
        if not bitstream.write_bit(bit):
            print_pretty("<Test bitstream.write_bit>", "Failed")
            fail += 1
    print_pretty("<Test bitstream.write_bit:out>", bitstream.pretty())

def test_read_bit(var_compare_bits):
    global fail
    print_pretty("<Test bitstream.read_bit:in>", bitstream.pretty())
    for bit in var_compare_bits:
        ok, bit2 = bitstream.read_bit()
        if not ok:
            print_pretty("<Test bitstream.read_bit>", "Failed (ok)")
            fail += 1
            continue
        if bit2 != bit:
            print_pretty("<Test bitstream.read_bit>", "Failed (match)")
            fail += 1
            continue
    print_pretty("<Test bitstream.read_bit:out>", bitstream.pretty())

def test_write(var_bytes):
    global fail
    print_pretty("<Test bitstream.write:in>", bitstream.pretty())
    if not bitstream.write(var_bytes, len(var_bytes) << 3):
        print_pretty("<Test bitstream.write>", "Failed")
        fail += 1
    print_pretty("<Test bitstream.write:out>", bitstream.pretty())

def test_read(var_compare_bytes):
    global fail
    print_pretty("<Test bitstream.read:in>", bitstream.pretty())
    ok, bytez = bitstream.read(len(var_compare_bytes) << 3)
    if ok:
        print(var_compare_bytes, bytez)
        if bytez != var_compare_bytes:
            print_pretty("<Test bitstream.read>", "Failed (match)")
            fail += 1
    else:
        print_pretty("<Test bitstream.read>", "Failed (ok)")
        fail += 1
    print_pretty("<Test bitstream.read:out>", bitstream.pretty())

def test_rw_bit(offset):
    bitstream._offset = offset
    test_write_bit(MAGIC_BITS)
    bitstream._offset = offset
    test_read_bit(MAGIC_BITS)

def test_rw_bytes(offset):
    bitstream._offset = offset
    test_write(MAGIC_BYTES)
    bitstream._offset = offset
    test_read(MAGIC_BYTES)

# Test all edges
for i in range(9):
    reset_bitstream()
    test_rw_bit(i)

reset_bitstream()

for i in range(9):
    reset_bitstream()
    test_rw_bytes(i)

print("Failed: %d" % fail)