from typing import (
    Iterable,
    Optional,
    Tuple,
)

from trie.exceptions import (
    PerfectVisibility,
)
from trie.fog import (
    HexaryTrieFog,
    TrieFrontierCache,
)
from trie.hexary import (
    HexaryTrie,
)
from trie.typing import (
    HexaryTrieNode,
    Nibbles,
)
from trie.utils.nibbles import (
    bytes_to_nibbles,
    nibbles_to_bytes,
)
from trie.utils.nodes import (
    consume_common_prefix,
)


class NodeIterator:
    """Iterate over all nodes of a trie, ensuring its consistency."""

    def __init__(self, trie: HexaryTrie) -> None:
        self._trie = trie

    def next(self, key_bytes: Optional[bytes] = None) -> Optional[bytes]:
        """
        Find the next key to the right from the given key, or None if there is
        no key to the right.

        .. NOTE:: To iterate the full trie, consider using keys() instead, for performance

        :param key_bytes: the key to start your search from. If None, return
            the first possible key.

        :return: key in bytes to the right of key_bytes, or None
        """
        root = self._trie.root_node
        none_traversed = Nibbles(())

        if key_bytes is None:
            next_key = self._get_next_key(root, none_traversed)
        else:
            key = bytes_to_nibbles(key_bytes)
            next_key = self._get_key_after(root, key, none_traversed)

        if next_key is None:
            return None
        else:
            return nibbles_to_bytes(next_key)

    def _get_key_after(
            self,
            node: HexaryTrieNode,
            key: Nibbles,
            traversed: Nibbles) -> Optional[Nibbles]:
        """
        Find the next key in the trie after key

        :param node: the source node to search for the next key after `key`
        :param key: the starting key used to seek the nearest key on the right
        :param traversed: the nibbles already traversed to get down to `node`

        :return: the complete key that is immediately to the right of `key` or None,
            if no key is immediately to the right (under `node`)
        """
        for next_segment in node.sub_segments:
            if key[:len(next_segment)] > next_segment:
                # This segment is to the left of the key, keep looking...
                continue
            else:
                # Either: found the exact match, or the next result to the right
                # Either way, we'll want to take a look
                next_node = self._trie.traverse_from(node, next_segment)

                common, key_remaining, segment_remaining = consume_common_prefix(key, next_segment)
                if len(segment_remaining) == 0:
                    # Found a perfect match! Keep looking for keys to the right of the target
                    next_key = self._get_key_after(
                        next_node,
                        key_remaining,
                        traversed + next_segment,
                    )
                    if next_key is None:
                        # Could not find a key to the right in any sub-node.
                        # In other words, *only* the target key is in this sub-trie
                        # So keep looking to the right...
                        continue
                    else:
                        # We successfully found a key to the right in a subtree, return it up
                        return next_key
                else:
                    # Found no exact match, and are now looking for the next possible key
                    return self._get_next_key(next_node, traversed + next_segment)

        if node.suffix > key:
            # This leaf node is to the right of the target key
            return traversed + node.suffix
        else:
            # Nothing found in any sub-segments
            return None

    def _get_next_key(self, node: HexaryTrieNode, traversed: Nibbles) -> Optional[Nibbles]:
        """
        Find the next possible key within the given node

        :param node: the parent node to search (plus all of its children)
        :param traversed: the key used to traverse down to `node`

        :return: the complete key that is the furthest left within `node`
        """
        if node.value:
            # This is either a leaf node, or a branch node with a value.
            # The value in a branch node comes before all the child values
            return traversed + node.suffix
        elif len(node.sub_segments) == 0:
            # Only leaves should have 0 sub-segments, and should have a value.
            # There shouldn't be any way to navigate to a blank node, as long as
            # the trie hasn't changed during iteration. If it has... I guess return None here.
            return None
        else:
            # This is a branch node with no value, or an extension node.
            # Either way, take the left-most child and repeat the search within it
            next_segment = node.sub_segments[0]
            next_node = self._trie.traverse_from(node, next_segment)
            return self._get_next_key(next_node, traversed + next_segment)

    def keys(self) -> Iterable[bytes]:
        """
        Iterate over all trie keys from left to right. Some performance benefit over
        using :meth:`next` repeatedly, by caching node accesses between yielded values.
        """
        for key, _ in self.items():
            yield key

    def items(self) -> Iterable[Tuple[bytes, bytes]]:
        """
        Iterate over all (key, value) pairs from left to right.
        """
        for prefix, node in self.nodes():
            if node.value:
                full_key = prefix + node.suffix
                yield nibbles_to_bytes(full_key), node.value

    def values(self) -> Iterable[bytes]:
        """
        Iterate over all stored values from left to right.
        """
        for _, node in self.nodes():
            if node.value:
                yield node.value

    def nodes(self) -> Iterable[Tuple[Nibbles, HexaryTrieNode]]:
        """
        Iterate over all trie nodes, starting at the left-most available one (the root),
        then the left-most available one (its left-most child) and so on.
        """
        next_fog = HexaryTrieFog()
        cache = TrieFrontierCache()

        while True:
            try:
                # Always get the furthest left unexplored value
                nearest_prefix = next_fog.nearest_right(())
            except PerfectVisibility:
                # No more unexplored nodes
                return

            try:
                cached_node, uncached_key = cache.get(nearest_prefix)
            except KeyError:
                node = self._trie.traverse(nearest_prefix)
            else:
                node = self._trie.traverse_from(cached_node, uncached_key)

            next_fog = next_fog.explore(nearest_prefix, node.sub_segments)

            if node.sub_segments:
                cache.add(nearest_prefix, node, node.sub_segments)
            else:
                cache.delete(nearest_prefix)

            yield nearest_prefix, node

from typing import (
    Iterable,
    Optional,
    Tuple,
)

from trie.exceptions import (
    PerfectVisibility,
)
from trie.fog import (
    HexaryTrieFog,
    TrieFrontierCache,
)
from trie.hexary import (
    HexaryTrie,
)
from trie.typing import (
    HexaryTrieNode,
    Nibbles,
)
from trie.utils.nibbles import (
    bytes_to_nibbles,
    nibbles_to_bytes,
)
from trie.utils.nodes import (
    consume_common_prefix,
)


class NodeIterator:
    """Iterate over all nodes of a trie, ensuring its consistency."""

    def __init__(self, trie: HexaryTrie) -> None:
        self._trie = trie

    def next(self, key_bytes: Optional[bytes] = None) -> Optional[bytes]:
        """
        Find the next key to the right from the given key, or None if there is
        no key to the right.

        .. NOTE:: To iterate the full trie, consider using keys() instead, for performance

        :param key_bytes: the key to start your search from. If None, return
            the first possible key.

        :return: key in bytes to the right of key_bytes, or None
        """
        root = self._trie.root_node
        none_traversed = Nibbles(())

        if key_bytes is None:
            next_key = self._get_next_key(root, none_traversed)
        else:
            key = bytes_to_nibbles(key_bytes)
            next_key = self._get_key_after(root, key, none_traversed)

        if next_key is None:
            return None
        else:
            return nibbles_to_bytes(next_key)

    def _get_key_after(
            self,
            node: HexaryTrieNode,
            key: Nibbles,
            traversed: Nibbles) -> Optional[Nibbles]:
        """
        Find the next key in the trie after key

        :param node: the source node to search for the next key after `key`
        :param key: the starting key used to seek the nearest key on the right
        :param traversed: the nibbles already traversed to get down to `node`

        :return: the complete key that is immediately to the right of `key` or None,
            if no key is immediately to the right (under `node`)
        """
        for next_segment in node.sub_segments:
            if key[:len(next_segment)] > next_segment:
                # This segment is to the left of the key, keep looking...
                continue
            else:
                # Either: found the exact match, or the next result to the right
                # Either way, we'll want to take a look
                next_node = self._trie.traverse_from(node, next_segment)

                common, key_remaining, segment_remaining = consume_common_prefix(key, next_segment)
                if len(segment_remaining) == 0:
                    # Found a perfect match! Keep looking for keys to the right of the target
                    next_key = self._get_key_after(
                        next_node,
                        key_remaining,
                        traversed + next_segment,
                    )
                    if next_key is None:
                        # Could not find a key to the right in any sub-node.
                        # In other words, *only* the target key is in this sub-trie
                        # So keep looking to the right...
                        continue
                    else:
                        # We successfully found a key to the right in a subtree, return it up
                        return next_key
                else:
                    # Found no exact match, and are now looking for the next possible key
                    return self._get_next_key(next_node, traversed + next_segment)

        if node.suffix > key:
            # This leaf node is to the right of the target key
            return traversed + node.suffix
        else:
            # Nothing found in any sub-segments
            return None

    def _get_next_key(self, node: HexaryTrieNode, traversed: Nibbles) -> Optional[Nibbles]:
        """
        Find the next possible key within the given node

        :param node: the parent node to search (plus all of its children)
        :param traversed: the key used to traverse down to `node`

        :return: the complete key that is the furthest left within `node`
        """
        if node.value:
            # This is either a leaf node, or a branch node with a value.
            # The value in a branch node comes before all the child values
            return traversed + node.suffix
        elif len(node.sub_segments) == 0:
            # Only leaves should have 0 sub-segments, and should have a value.
            # There shouldn't be any way to navigate to a blank node, as long as
            # the trie hasn't changed during iteration. If it has... I guess return None here.
            return None
        else:
            # This is a branch node with no value, or an extension node.
            # Either way, take the left-most child and repeat the search within it
            next_segment = node.sub_segments[0]
            next_node = self._trie.traverse_from(node, next_segment)
            return self._get_next_key(next_node, traversed + next_segment)

    def keys(self) -> Iterable[bytes]:
        """
        Iterate over all trie keys from left to right. Some performance benefit over
        using :meth:`next` repeatedly, by caching node accesses between yielded values.
        """
        for key, _ in self.items():
            yield key

    def items(self) -> Iterable[Tuple[bytes, bytes]]:
        """
        Iterate over all (key, value) pairs from left to right.
        """
        for prefix, node in self.nodes():
            if node.value:
                full_key = prefix + node.suffix
                yield nibbles_to_bytes(full_key), node.value

    def values(self) -> Iterable[bytes]:
        """
        Iterate over all stored values from left to right.
        """
        for _, node in self.nodes():
            if node.value:
                yield node.value

    def nodes(self) -> Iterable[Tuple[Nibbles, HexaryTrieNode]]:
        """
        Iterate over all trie nodes, starting at the left-most available one (the root),
        then the left-most available one (its left-most child) and so on.
        """
        next_fog = HexaryTrieFog()
        cache = TrieFrontierCache()

        while True:
            try:
                # Always get the furthest left unexplored value
                nearest_prefix = next_fog.nearest_right(())
            except PerfectVisibility:
                # No more unexplored nodes
                return

            try:
                cached_node, uncached_key = cache.get(nearest_prefix)
            except KeyError:
                node = self._trie.traverse(nearest_prefix)
            else:
                node = self._trie.traverse_from(cached_node, uncached_key)

            next_fog = next_fog.explore(nearest_prefix, node.sub_segments)

            if node.sub_segments:
                cache.add(nearest_prefix, node, node.sub_segments)
            else:
                cache.delete(nearest_prefix)

            yield nearest_prefix, node

Iterate over all nodes of a trie, ensuring its consistency. Find the next key to the right from the given key, or None if there is         no key to the right.          .. NOTE:: To iterate the full trie, consider using keys() instead, for performance          :param key_bytes: the key to start your search from. If None, return             the first possible key.          :return: key in bytes to the right of key_bytes, or None Find the next key in the trie after key          :param node: the source node to search for the next key after `key`         :param key: the starting key used to seek the nearest key on the right         :param traversed: the nibbles already traversed to get down to `node`          :return: the complete key that is immediately to the right of `key` or None,             if no key is immediately to the right (under `node`) # This segment is to the left of the key, keep looking... # Either: found the exact match, or the next result to the right # Either way, we'll want to take a look # Found a perfect match! Keep looking for keys to the right of the target # Could not find a key to the right in any sub-node. # In other words, *only* the target key is in this sub-trie # So keep looking to the right... # We successfully found a key to the right in a subtree, return it up # Found no exact match, and are now looking for the next possible key # This leaf node is to the right of the target key # Nothing found in any sub-segments Find the next possible key within the given node          :param node: the parent node to search (plus all of its children)         :param traversed: the key used to traverse down to `node`          :return: the complete key that is the furthest left within `node` # This is either a leaf node, or a branch node with a value. # The value in a branch node comes before all the child values # Only leaves should have 0 sub-segments, and should have a value. # There shouldn't be any way to navigate to a blank node, as long as # the trie hasn't changed during iteration. If it has... I guess return None here. # This is a branch node with no value, or an extension node. # Either way, take the left-most child and repeat the search within it Iterate over all trie keys from left to right. Some performance benefit over         using :meth:`next` repeatedly, by caching node accesses between yielded values. Iterate over all (key, value) pairs from left to right. Iterate over all stored values from left to right. Iterate over all trie nodes, starting at the left-most available one (the root),         then the left-most available one (its left-most child) and so on. # Always get the furthest left unexplored value # No more unexplored nodes