from typing import TypeVar, Generic, Callable, MutableMapping, List, Iterable, Union, Any

K = TypeVar('K')
V = TypeVar('V')

KeyCallable = Callable[[V], Union[K, Iterable[K]]]

__all__ = [
    "Index",
    "ListIndex",
    "UniqueIndex",
    "ObjDb",
]


class Index(Generic[K, V]):
    def add(self, value: V):
        pass

    def get(self, key: K) -> Iterable[V]:
        pass

    def clear(self):
        pass

    def __iter__(self):
        pass

    def items(self):
        pass

    def values(self):
        pass


class ListIndex(Index[K, V]):
    def __init__(self, key_callable: KeyCallable):
        self.idx: MutableMapping[K, V] = {}
        self.key_callable = key_callable

    def add(self, value: V):
        keys = self.key_callable(value)

        if keys is None:
            return

        for key in keys:
            if key is None:
                continue

            values = self.idx.get(key, None)
            if values is not None:
                values.append(value)
            else:
                self.idx[key] = [value]

    def get(self, key: K) -> Iterable[V]:
        return [self.idx[key]]

    def clear(self):
        return self.idx.clear()

    def __iter__(self):
        return self.idx.__iter__()

    def items(self):
        return self.idx.items()

    def values(self):
        return self.idx.values()


class UniqueIndex(Index[K, V]):
    def __init__(self, key_callable: KeyCallable):
        self.idx: MutableMapping[K, V] = {}
        self.key_callable = key_callable

    def get(self, key: K) -> Iterable[V]:
        return [self.idx[key]]

    def get_single(self, key: K) -> V:
        return self.idx[key]

    def add(self, value: V):
        keys = self.key_callable(value)

        if not isinstance(keys, Iterable):
            keys = [keys]

        for key in keys:
            key = self.key_callable(value)
            present = self.idx.get(key, None)
            if present is not None:
                raise KeyError("Duplicate part in index: {}".format(key))

            self.idx[key] = True

    def clear(self):
        return self.idx.clear()

    def __iter__(self):
        return self.idx.__iter__()

    def items(self):
        return self.idx.items()

    def values(self):
        return self.idx.values()


class ObjDb(Generic[V]):
    def __init__(self):
        self.values: List[V] = []
        self._indexes: MutableMapping[str, Index[Any, V]] = {}

    def add(self, value: V):
        for idx in self._indexes.values():
            idx.add(value)

        self.values.append(value)

    def __iter__(self):
        return self.values.__iter__()

    def __len__(self) -> int:
        return len(self.values)

    def add_index(self, name, key_callable: KeyCallable) -> ListIndex[Any, V]:
        idx = ListIndex(key_callable)
        self._indexes[name] = idx

        for value in self.values:
            idx.add(value)

        return idx

    def add_unique_index(self, name, key_callable: KeyCallable) -> UniqueIndex[Any, V]:
        idx = UniqueIndex(key_callable)
        self._indexes[name] = idx
        return idx

    def index(self, name) -> Index:
        return self._indexes[name]