from typing import Generic, Iterable, TypeVar

from rainbowadn.core import *
from rainbowadn.nullability import *

__all__ = ('BinaryTree', 'BinaryTreeFactory',)

TreeKeyType = TypeVar('TreeKeyType')


class BinaryTree(RecursiveMentionable, Generic[TreeKeyType]):
    def __factory__(self) -> RainbowFactory['BinaryTree[TreeKeyType]']:
        return self.factory(self.key.__factory__())

    @classmethod
    def factory(cls, factory: RainbowFactory[TreeKeyType]) -> RainbowFactory['BinaryTree[TreeKeyType]']:
        assert isinstance(factory, RainbowFactory)
        return BinaryTreeFactory(factory)

    def __init__(
            self,
            treel: NullableReference['BinaryTree[TreeKeyType]'],
            treer: NullableReference['BinaryTree[TreeKeyType]'],
            key: TreeKeyType
    ):
        assert isinstance(treel, NullableReference)
        assert isinstance(treer, NullableReference)
        assert isinstance(key, Mentionable)
        self.treel = treel
        self.treer = treer
        self.key = key

    def points(self) -> Iterable[HashPoint]:
        return [*self.treel.points(), *self.treer.points(), *RecursiveMentionable.points_of(self.key)]

    def __bytes__(self):
        return bytes(self.treel) + bytes(self.treer) + bytes(self.key)

    async def str(self, tab: int) -> str:
        assert isinstance(tab, int)
        treel_str, key_str, treer_str = await gather(
            self.treel.str(tab),
            hash_point_format(HashPoint.of(self.key), tab),
            self.treer.str(tab),
        )
        assert isinstance(treel_str, str)
        assert isinstance(key_str, str)
        assert isinstance(treer_str, str)
        return f'{treel_str}' \
               f'{tabulate(tab)}{key_str}' \
               f'{tabulate(tab)}{treer_str}'


class BinaryTreeFactory(RainbowFactory[BinaryTree[TreeKeyType]], Generic[TreeKeyType]):
    def __init__(self, factory: RainbowFactory[TreeKeyType]):
        assert isinstance(factory, RainbowFactory)
        self.factory = factory
        self.reference_factory = NullableReference.f(self)

    def from_bytes(self, source: bytes, resolver: HashResolver) -> BinaryTree[TreeKeyType]:
        assert isinstance(source, bytes)
        assert isinstance(resolver, HashResolver)
        return BinaryTree(
            self.reference_factory.from_bytes(source[:HashPoint.HASH_LENGTH], resolver),
            self.reference_factory.from_bytes(
                source[HashPoint.HASH_LENGTH:HashPoint.HASH_LENGTH * 2],
                resolver
            ),
            self.factory.from_bytes(source[HashPoint.HASH_LENGTH * 2:], resolver),
        )

    def loose(self) -> RainbowFactory[BinaryTree[TreeKeyType]]:
        return self