import numpy as np
from towhee import register
from pymilvus import Collection


@register(output_schema=['mr'])
class Milvus:
    """
    Milvus ANN index class.
    """

    def __init__(self, collection):
        """
        Get an existing collection.
        """
        if isinstance(collection, str):
            collection = Collection(collection)
        self._collection = collection

    def __call__(self, data):
        """
        Insert data to Milvus.

        Args:
        data (`list`):
            The data to insert into milvus.

        Returns:
            A MutationResult object contains `insert_count` represents how many and a `primary_keys` of primary keys.

        """
        vectors = []
        if isinstance(data, np.ndarray):
            data = [data]
        for v in data:
            vectors.append(v if isinstance(v, list) else [v])
        mr = self._collection.insert(vectors)
        return mr