import typing

from chainable_iterator import ChainableIterator

Key, Value = typing.Any, typing.Any
KeyPath = typing.List[Key]
Predicate = typing.Callable[[KeyPath, Key, Value], bool]
_Node = typing.NamedTuple('Node', (('key_path', KeyPath), ('data', typing.Dict)))


def _match_all(_p: KeyPath, _k: Key, _v: Value) \
        -> bool:
    return True


def _decompose_dict(the_dict: typing.Dict, predicate: Predicate = _match_all) \
        -> typing.Generator[typing.Tuple[KeyPath, Value], None, None]:
    root = ChainableIterator((_Node([], the_dict),))
    for node in root:
        key_path = node.key_path
        for k, v in node.data.items():
            key_path.append(k)
            if isinstance(v, dict):
                root.chain((_Node(key_path, v),))
            if predicate(key_path, k, v):
                yield key_path, v
            key_path = key_path[:-1]


def _recompose_dict(items: typing.Iterator[typing.Tuple[KeyPath, Value]]) \
        -> typing.Dict:
    root = {}
    for path, value in items:
        base, last = root, len(path) - 1
        for i, key in enumerate(path):
            base = base.setdefault(key, {} if i < last else value)
    return root


def filter_dict(the_dict: typing.Dict, predicate: Predicate = _match_all) \
        -> typing.Dict:
    return _recompose_dict(_decompose_dict(the_dict, predicate))


if __name__ == '__main__':
    import json

    data = {
        '*': {
            'a': {
                'x': {
                    1: None,
                },
                'k': {'lorem': 'spam'},
                '-': [
                    {
                        'h': 1,
                    },
                    {
                        'o': {
                            '/1': 1,
                            '/3': 3
                        },
                    },
                    {
                        'h': 3,
                    }
                ]
            },
            'b': {
                'y': {
                    1: None,
                },
                'k': {'lorem': 'eggs'},
                '-': [
                    {
                        'h': 4,
                    },
                    {
                        '.': {
                            '/4': 4,
                            '/6': 6
                        },
                    },
                    {
                        'h': 6,
                    }
                ]
            },
            'c': {
                'z': {
                    1: None,
                },
                'k': {'lorem': 'ipsum'},
                '-': [
                    {
                        'h': 7,
                    },
                    {
                        'o': {
                            '/7': 7,
                            '/9': 9
                        },
                    },
                    {
                        'h': 9,
                    }
                ]
            }
        }
    }

    print(json.dumps(filter_dict(data, lambda p, k, v: v is None), indent=2))

    print()

    print(json.dumps(filter_dict(data, lambda p, k, v: k == 'lorem' and v == 'ipsum'), indent=2))

    print()

    print(json.dumps(
        filter_dict(data,
                    lambda p, k, v: isinstance(v, list) and any(map(
                        lambda d: filter_dict(d, lambda p1, k1, v1: k1 == '.'), v))),
        indent=2))