from optional import Optional
from remerkleable.basic import boolean, uint8, uint16, uint32, uint64, uint128, uint256
from remerkleable.bitfields import Bitlist, Bitvector
from remerkleable.complex import Container, List, Vector
from remerkleable.union import Union

def do_test(value):
    v = value.get()
    if v is None:
        assert value.encode_bytes() == b''
        assert value.hash_tree_root() == List[value.__class__, 1]().hash_tree_root()
    else:
        assert value.encode_bytes() == v.encode_bytes()
        assert value.hash_tree_root() == List[value.__class__, 1](v).hash_tree_root()
    assert value.__class__.decode_bytes(value.encode_bytes()) == value

if __name__ == '__main__':
    do_test(Optional[uint8](None))
    do_test(Optional[uint8](8))

    do_test(Optional[uint16](None))
    do_test(Optional[uint16](16))

    do_test(Optional[uint32](None))
    do_test(Optional[uint32](32))

    do_test(Optional[uint64](None))
    do_test(Optional[uint64](64))

    do_test(Optional[uint128](None))
    do_test(Optional[uint128](128))

    do_test(Optional[uint256](None))
    do_test(Optional[uint256](256))

    do_test(Optional[boolean](None))
    do_test(Optional[boolean](True))

    class Foo(Container):
        a: uint64
        b: Optional[uint32]
        c: Optional[uint16]

    do_test(Optional[Foo](None))
    do_test(Optional[Foo](Foo(a=64)))
    do_test(Optional[Foo](Foo(a=64, b=Optional[uint32](32))))
    do_test(Optional[Foo](Foo(a=64, b=Optional[uint32](32), c=Optional[uint16](16))))

    do_test(Optional[Vector[uint64, 1]](None))
    do_test(Optional[Vector[uint64, 1]](Vector[uint64, 1](64)))
    do_test(Optional[Vector[uint64, 5]](None))
    do_test(Optional[Vector[uint64, 5]](Vector[uint64, 5](64, 64, 64, 64, 64)))

    do_test(Optional[Bitvector[1]](None))
    do_test(Optional[Bitvector[1]](Bitvector[1](True)))
    do_test(Optional[Bitvector[9]](None))
    do_test(Optional[Bitvector[9]](Bitvector[9](True, True, True, True, False, True, True, True, True)))

    do_test(Optional[Bitlist[0]](None))
    do_test(Optional[Bitlist[0]](Bitlist[0]()))
    do_test(Optional[Bitlist[1]](None))
    do_test(Optional[Bitlist[1]](Bitlist[1](True)))
    do_test(Optional[Bitlist[9]](None))
    do_test(Optional[Bitlist[9]](Bitlist[9](True)))

    do_test(Optional[Union[None, uint64, uint32]](None))
    do_test(Optional[Union[None, uint64, uint32]](Union[None, uint64, uint32](selector=0)))
    do_test(Optional[Union[None, uint64, uint32]](Union[None, uint64, uint32](selector=1, value=64)))
    do_test(Optional[Union[None, uint64, uint32]](Union[None, uint64, uint32](selector=2, value=32)))