from collections import ChainMap
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    ItemsView,
    Iterable,
    MutableMapping,
    NamedTuple,
    Optional,
    Type,
    TypeVar,
    Union,
    ValuesView,
)

from docarray.array.list_advance_indexing import ListAdvancedIndexing
from docarray.typing import NdArray
from docarray.typing.tensor.abstract_tensor import AbstractTensor

if TYPE_CHECKING:
    from docarray.array.doc_vec.doc_vec import DocVec

IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]


T = TypeVar('T', bound='ColumnStorage')


class ColumnsJsonCompatible(NamedTuple):
    tensor_columns: Dict[str, Any]
    doc_columns: Dict[str, Any]
    docs_vec_columns: Dict[str, Any]
    any_columns: Dict[str, Any]


class ColumnStorage:
    """
    ColumnStorage is a container to store the columns of the
    :class:`~docarray.array.doc_vec.DocVec`.

    :param tensor_columns: a Dict of AbstractTensor
    :param doc_columns: a Dict of :class:`~docarray.array.doc_vec.DocVec`
    :param docs_vec_columns: a Dict of List of :class:`~docarray.array.doc_vec.DocVec`
    :param any_columns: a Dict of List
    :param tensor_type: Class used to wrap the doc_vec tensors
    """

    def __init__(
        self,
        tensor_columns: Dict[str, Optional[AbstractTensor]],
        doc_columns: Dict[str, Optional['DocVec']],
        docs_vec_columns: Dict[str, Optional[ListAdvancedIndexing['DocVec']]],
        any_columns: Dict[str, ListAdvancedIndexing],
        tensor_type: Type[AbstractTensor] = NdArray,
    ):
        self.tensor_columns = tensor_columns
        self.doc_columns = doc_columns
        self.docs_vec_columns = docs_vec_columns
        self.any_columns = any_columns

        self.tensor_type = tensor_type

        self.columns = ChainMap(  # type: ignore
            self.tensor_columns,  # type: ignore
            self.doc_columns,  # type: ignore
            self.docs_vec_columns,  # type: ignore
            self.any_columns,  # type: ignore
        )  # type: ignore

    def __len__(self) -> int:
        return len(self.any_columns['id'])  # TODO what if ID are None ?

    def __getitem__(self: T, item: IndexIterType) -> T:
        if isinstance(item, tuple):
            item = list(item)
        tensor_columns = {
            key: col[item] if col is not None else None
            for key, col in self.tensor_columns.items()
        }
        doc_columns = {
            key: col[item] if col is not None else None
            for key, col in self.doc_columns.items()
        }
        docs_vec_columns = {
            key: col[item] if col is not None else None
            for key, col in self.docs_vec_columns.items()
        }
        any_columns = {
            key: col[item] if col is not None else None
            for key, col in self.any_columns.items()
        }

        return self.__class__(
            tensor_columns,
            doc_columns,
            docs_vec_columns,
            any_columns,
            self.tensor_type,
        )

    def columns_json_compatible(self) -> ColumnsJsonCompatible:
        tens_cols = {
            key: value._docarray_to_json_compatible() if value is not None else value
            for key, value in self.tensor_columns.items()
        }
        doc_cols = {
            key: value._docarray_to_json_compatible() if value is not None else value
            for key, value in self.doc_columns.items()
        }
        doc_vec_cols = {
            key: [vec._docarray_to_json_compatible() for vec in value]
            if value is not None
            else value
            for key, value in self.docs_vec_columns.items()
        }
        return ColumnsJsonCompatible(
            tens_cols, doc_cols, doc_vec_cols, self.any_columns
        )

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, ColumnStorage):
            return False
        if self.tensor_type != other.tensor_type:
            return False
        for col_map_self, col_map_other in zip(self.columns.maps, other.columns.maps):
            if col_map_self.keys() != col_map_other.keys():
                return False
            for key_self in col_map_self.keys():
                if key_self == 'id':
                    continue

                val1, val2 = col_map_self[key_self], col_map_other[key_self]
                if isinstance(val1, AbstractTensor):
                    values_are_equal = val1.get_comp_backend().equal(val1, val2)
                elif isinstance(val2, AbstractTensor):
                    values_are_equal = val2.get_comp_backend().equal(val1, val2)
                else:
                    values_are_equal = val1 == val2
                if not values_are_equal:
                    return False
        return True


class ColumnStorageView(dict, MutableMapping[str, Any]):
    index: int
    storage: ColumnStorage

    def __init__(self, index: int, storage: ColumnStorage):
        super().__init__()
        self.index = index
        self.storage = storage

    def __getitem__(self, name: str) -> Any:
        if name in self.storage.tensor_columns.keys():
            tensor = self.storage.tensor_columns[name]
            if tensor is None:
                return None
            if tensor.get_comp_backend().n_dim(tensor) == 1:
                # to ensure consistensy between numpy and pytorch
                # we wrap the scalr in a tensor of ndim = 1
                # otherwise numpy pass by value whereas torch by reference
                col = self.storage.tensor_columns[name]

                if col is not None:
                    return col[self.index : self.index + 1]
                else:
                    return None

        col = self.storage.columns[name]

        if col is None:
            return None
        return col[self.index]

    def __reduce__(self):
        # implementing __reduce__ to solve a pickle issue when subclassing dict
        # see here: https://stackoverflow.com/questions/21144845/how-can-i-unpickle-a-subclass-of-dict-that-validates-with-setitem-in-pytho
        return (ColumnStorageView, (self.index, self.storage))

    def __setitem__(self, name, value) -> None:
        if self.storage.columns[name] is None:
            raise ValueError(
                f'Cannot set an item to a None column. This mean that '
                f'the DocVec that encapsulate this doc has the field '
                f'{name} set to None. If you want to modify that you need to do it at the'
                f'DocVec level. `docs.field = np.zeros(10)`'
            )

        self.storage.columns[name][self.index] = value

    def __delitem__(self, key):
        raise RuntimeError('Cannot delete an item from a StorageView')

    def __iter__(self):
        return self.storage.columns.keys()

    def __len__(self):
        return len(self.storage.columns)

    def _local_dict(self):
        """The storage.columns dictionary with every value at position self.index"""

        return {key: self[key] for key in self.storage.columns.keys()}

    def keys(self):
        return self.storage.columns.keys()

    # type ignore because return type dict_values is private and we cannot use it.
    # context: https://github.com/python/typing/discussions/1033
    def values(self) -> ValuesView:  # type: ignore
        return ValuesView(self._local_dict())

    # type ignore because return type dict_items is private and we cannot use it.
    # context: https://github.com/python/typing/discussions/1033
    def items(self) -> ItemsView:  # type: ignore
        return ItemsView(self._local_dict())

    def to_dict(self) -> Dict[str, Any]:
        """
        Return a dictionary with the same keys as the storage.columns
        and the values at position self.index.
        Warning: modification on the dict will not be reflected on the storage.
        """
        return {key: self[key] for key in self.storage.columns.keys()}
