import copy
import uuid
from typing import Optional, TYPE_CHECKING, Union, Dict, Iterable, List, Tuple
from dataclasses import dataclass, field
import re

import numpy as np
from pymilvus import (
    connections,
    Collection,
    FieldSchema,
    DataType,
    CollectionSchema,
    has_collection,
    loading_progress,
)

from docarray import Document, DocumentArray
from docarray.array.storage.base.backend import BaseBackendMixin, TypeMap
from docarray.helper import dataclass_from_dict, _safe_cast_int
from docarray.score import NamedScore

if TYPE_CHECKING:
    from docarray.typing import (
        DocumentArraySourceType,
    )


ID_VARCHAR_LEN = 1024
SERIALIZED_VARCHAR_LEN = (
    65_535  # 65_535 is the maximum that Milvus allows for a VARCHAR field
)
COLUMN_VARCHAR_LEN = 1024
OFFSET_VARCHAR_LEN = 1024


def _always_true_expr(primary_key: str) -> str:
    """
    Returns a Milvus expression that is always true, thus allowing for the retrieval of all entries in a Collection
    Assumes that the primary key is of type DataType.VARCHAR

    :param primary_key: the name of the primary key
    :return: a Milvus expression that is always true for that primary key
    """
    return f'({primary_key} in ["1"]) or ({primary_key} not in ["1"])'


def _ids_to_milvus_expr(ids):
    ids = ['"' + _id + '"' for _id in ids]
    return '[' + ','.join(ids) + ']'


def _batch_list(l: List, batch_size: int):
    """Iterates over a list in batches of size batch_size"""
    if batch_size < 1:
        yield l
        return
    l_len = len(l)
    for ndx in range(0, l_len, batch_size):
        yield l[ndx : min(ndx + batch_size, l_len)]


def _sanitize_collection_name(name):
    """Removes all chars that are not allowed in a Milvus collection name.
    Thus, it removes all chars that are not alphanumeric or an underscore.

    :param name: the collection name to sanitize
    :return: the sanitized collection name.
    """
    return ''.join(
        re.findall('[a-zA-Z0-9_]', name)
    )  # remove everything that is not a letter, number or underscore


@dataclass
class MilvusConfig:
    n_dim: int
    collection_name: str = None
    host: str = 'localhost'
    port: Optional[Union[str, int]] = 19530  # 19530 for gRPC, 9091 for HTTP
    distance: str = 'IP'  # metric_type in milvus
    index_type: str = 'HNSW'
    index_params: Dict = field(
        default_factory=lambda: {
            'M': 4,
            'efConstruction': 200,
        }
    )  # passed to milvus at index creation time. The default assumes 'HNSW' index type
    collection_config: Dict = field(
        default_factory=dict
    )  # passed to milvus at collection creation time
    serialize_config: Dict = field(default_factory=dict)
    consistency_level: str = 'Session'
    batch_size: int = -1
    columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
    list_like: bool = True
    root_id: bool = True


class BackendMixin(BaseBackendMixin):

    TYPE_MAP = {
        'str': TypeMap(type=DataType.VARCHAR, converter=str),
        'float': TypeMap(
            type=DataType.DOUBLE, converter=float
        ),  # it doesn't like DataType.FLOAT type, perhaps because python floats are double precision?
        'double': TypeMap(type=DataType.DOUBLE, converter=float),
        'int': TypeMap(type=DataType.INT64, converter=_safe_cast_int),
        'bool': TypeMap(type=DataType.BOOL, converter=bool),
    }

    def _init_storage(
        self,
        _docs: Optional['DocumentArraySourceType'] = None,
        config: Optional[Union[MilvusConfig, Dict]] = None,
        **kwargs,
    ):
        config = copy.deepcopy(config)
        if not config:
            raise ValueError('Empty config is not allowed for Milvus storage')
        elif isinstance(config, dict):
            config = dataclass_from_dict(MilvusConfig, config)

        if config.collection_name is None:
            id = uuid.uuid4().hex
            config.collection_name = 'docarray__' + id
        self._list_like = config.list_like
        self._config = config
        self._config.columns = self._normalize_columns(self._config.columns)

        self._connection_alias = f'docarray_{config.host}_{config.port}'
        connections.connect(
            alias=self._connection_alias, host=config.host, port=config.port
        )

        self._collection = self._create_or_reuse_collection()
        self._offset2id_collection = self._create_or_reuse_offset2id_collection()
        self._build_index()
        super()._init_storage(**kwargs)

        # To align with Sqlite behavior; if `docs` is not `None` and table name
        # is provided, :class:`DocumentArraySqlite` will clear the existing
        # table and load the given `docs`
        if _docs is None:
            return

        self.clear()
        if isinstance(_docs, Iterable):
            self.extend(_docs)
        else:
            if isinstance(_docs, Document):
                self.append(_docs)

    def _create_or_reuse_collection(self):
        if has_collection(self._config.collection_name, using=self._connection_alias):
            return Collection(
                self._config.collection_name, using=self._connection_alias
            )

        document_id = FieldSchema(
            name='document_id',
            dtype=DataType.VARCHAR,
            max_length=ID_VARCHAR_LEN,
            is_primary=True,
        )
        embedding = FieldSchema(
            name='embedding', dtype=DataType.FLOAT_VECTOR, dim=self._config.n_dim
        )
        serialized = FieldSchema(
            name='serialized', dtype=DataType.VARCHAR, max_length=SERIALIZED_VARCHAR_LEN
        )

        additional_columns = []
        for col, coltype in self._config.columns.items():
            mapped_type = self._map_type(coltype)
            if mapped_type == DataType.VARCHAR:
                field_ = FieldSchema(
                    name=col, dtype=mapped_type, max_length=COLUMN_VARCHAR_LEN
                )
            else:
                field_ = FieldSchema(name=col, dtype=mapped_type)
            additional_columns.append(field_)

        schema = CollectionSchema(
            fields=[document_id, embedding, serialized, *additional_columns],
            description='DocumentArray collection schema',
        )
        return Collection(
            name=self._config.collection_name,
            schema=schema,
            using=self._connection_alias,
            **self._config.collection_config,
        )

    def _build_index(self):
        index_params = {
            'metric_type': self._config.distance,
            'index_type': self._config.index_type,
            'params': self._config.index_params,
        }
        self._collection.create_index(field_name='embedding', index_params=index_params)

    def _create_or_reuse_offset2id_collection(self):
        if has_collection(
            self._config.collection_name + '_offset2id', using=self._connection_alias
        ):
            return Collection(
                self._config.collection_name + '_offset2id',
                using=self._connection_alias,
            )

        document_id = FieldSchema(
            name='document_id', dtype=DataType.VARCHAR, max_length=ID_VARCHAR_LEN
        )
        offset = FieldSchema(
            name='offset',
            dtype=DataType.VARCHAR,
            max_length=OFFSET_VARCHAR_LEN,
            is_primary=True,
        )
        dummy_vector = FieldSchema(
            name='dummy_vector', dtype=DataType.FLOAT_VECTOR, dim=1
        )

        schema = CollectionSchema(
            fields=[offset, document_id, dummy_vector],
            description='offset2id for DocumentArray',
        )

        return Collection(
            name=self._config.collection_name + '_offset2id',
            schema=schema,
            using=self._connection_alias,
            # **self._config.collection_config,  # we probably don't want to apply the same config here
        )

    def _ensure_unique_config(
        self,
        config_root: dict,
        config_subindex: dict,
        config_joined: dict,
        subindex_name: str,
    ) -> dict:
        if 'collection_name' not in config_subindex:
            config_joined['collection_name'] = _sanitize_collection_name(
                config_joined['collection_name'] + '_subindex_' + subindex_name
            )
        return config_joined

    def _doc_to_milvus_payload(self, doc):
        return self._docs_to_milvus_payload([doc])

    def _docs_to_milvus_payload(self, docs: 'Iterable[Document]'):
        extra_columns = [
            [self._map_column(doc.tags.get(col), col_type) for doc in docs]
            for col, col_type in self._config.columns.items()
        ]
        return [
            [doc.id for doc in docs],
            [self._map_embedding(doc.embedding) for doc in docs],
            [doc.to_base64(**self._config.serialize_config) for doc in docs],
            *extra_columns,
        ]

    @staticmethod
    def _docs_from_query_response(response):
        return DocumentArray([Document.from_base64(d['serialized']) for d in response])

    @staticmethod
    def _docs_from_search_response(responses, distance: str) -> 'List[DocumentArray]':
        das = []
        for r in responses:
            da = []
            for hit in r:
                doc = Document.from_base64(hit.entity.get('serialized'))
                doc.scores[distance] = NamedScore(value=hit.score)
                da.append(doc)
            das.append(DocumentArray(da))
        return das

    def _update_kwargs_from_config(self, field_to_update, **kwargs):
        kwargs_field_value = kwargs.get(field_to_update, None)
        config_field_value = getattr(self._config, field_to_update, None)

        if (
            kwargs_field_value is not None or config_field_value is None
        ):  # no need to update
            return kwargs

        kwargs[field_to_update] = config_field_value
        return kwargs

    def _map_embedding(self, embedding):
        if embedding is not None:
            from docarray.math.ndarray import to_numpy_array

            embedding = to_numpy_array(embedding)

            if embedding.ndim > 1:
                embedding = np.asarray(embedding).squeeze()
        else:
            embedding = np.zeros(self._config.n_dim)
        return embedding

    def __getstate__(self):
        d = dict(self.__dict__)
        del d['_collection']
        del d['_offset2id_collection']
        return d

    def __setstate__(self, state):
        self.__dict__ = state
        connections.connect(
            alias=self._connection_alias, host=self._config.host, port=self._config.port
        )
        self._collection = self._create_or_reuse_collection()
        self._offset2id_collection = self._create_or_reuse_offset2id_collection()

    def __enter__(self):
        _ = super().__enter__()
        self._collection.load()
        self._offset2id_collection.load()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._collection.release()
        self._offset2id_collection.release()
        super().__exit__(exc_type, exc_val, exc_tb)

    def loaded_collection(self, collection=None):
        """
        Context manager to load a collection and release it after the context is exited.
        If the collection is already loaded when entering, it will not be released while exiting.

        :param collection: the collection to load. If None, the main collection of this indexer is used.
        :return: Context manager for the provided collection.
        """

        class LoadedCollectionManager:
            def __init__(self, coll, connection_alias):
                self._collection = coll
                self._loaded_when_enter = False
                self._connection_alias = connection_alias

            def __enter__(self):
                self._loaded_when_enter = (
                    loading_progress(
                        self._collection.name, using=self._connection_alias
                    )['loading_progress']
                    != '0%'
                )
                self._collection.load()
                return self

            def __exit__(self, exc_type, exc_val, exc_tb):
                if not self._loaded_when_enter:
                    self._collection.release()

        return LoadedCollectionManager(
            collection if collection else self._collection, self._connection_alias
        )
