from pymilvusdm.core.milvus_client import MilvusIndex
import math
import time
import sys


class DataToMilvus:
    def __init__(self, logger, client):
        self.client = client
        self.logger = logger
        self.batch_size = 100000

    def if_create_collection(self, collection_name, partition_tag):
        try:
            if self.client.has_collection(collection_name):
                if not partition_tag:
                    self.logger.info("There collection:{} exists, and skip it.".format(collection_name))
                    sys.exit(1)
                elif partition_tag and self.client.has_partition(collection_name, partition_tag):
                    self.logger.info("There collection:{} partition:{} exists, and skip it.".format(collection_name, partition_tag))
                    sys.exit(1)
        except Exception as e:
            self.logger.error('Error with: {}'.format(e))
            sys.exit(1)

    def creat_none_partition(self, collection, collection_param, mode, has_collection):
        _continue = True
        if not has_collection:
            self.client.create_collection(collection, collection_param)
        else:
            if mode == 'overwrite':
                self.client.drop_collection(collection)
                while has_collection:
                    has_collection = self.client.has_collection(collection)
                self.client.create_collection(collection, collection_param)
            elif mode == 'skip':
                _continue = False
        return _continue

    def creat_assigned_partition(self, collection, collection_param, mode, partition, has_collection):
        _continue = True
        if not has_collection:
            self.client.create_collection(collection, collection_param)
            self.client.creat_partition(collection, partition)
        else:
            has_partition = self.client.has_partition(collection, partition)
            if not has_partition:
                self.client.creat_partition(collection, partition)
            else:
                if mode == 'overwrite':
                    self.client.drop_partition(collection, partition)
                    while has_partition:
                        has_partition = self.client.has_partition(collection, partition)
                    self.client.creat_partition(collection, partition)
                elif mode == 'skip':
                    _continue = False
        return _continue

    def insert_data(self, vectors, collection, collection_param_1, mode, ids=None, partition=None):
        try:
            data_type = self.client.get_data_type(collection_param_1)
            collection_param = {'collection_name': collection, 'dimension': collection_param_1['dimension'],
                                'index_file_size': collection_param_1['index_file_size'], 'data_type': data_type}
            has_collection = self.client.has_collection(collection)
            if partition:
                _continue = self.creat_assigned_partition(collection, collection_param, mode, partition, has_collection)
            else:
                _continue = self.creat_none_partition(collection, collection_param, mode, has_collection)
            return_ids = []
            if _continue:
                for i in range(math.ceil(len(vectors) / self.batch_size)):
                    _ids = None
                    if ids.size != 0:
                        _ids = ids[i * self.batch_size: (i + 1) * self.batch_size]
                    vector = vectors[i * self.batch_size: (i + 1) * self.batch_size]
                    if not isinstance(vector, list):
                        vector = vector.tolist()
                    status, ids_ = self.client.insert(collection, vector, _ids.tolist(), partition)
                    return_ids += ids_

                if len(vectors) == len(return_ids):
                    self.logger.debug(
                        "Successfuly insert collection: {}/partition: {} , total num: {}".format(
                            collection, partition, len(return_ids)))
                else:
                    raise Exception(
                        "ERROR: The collection: {}/partition: {} data count is not equal!".format(collection,
                                                                                                  partition))
                    # self.logger.error(
                    #     "ERROR: The collection: {}/partition: {} data count is not equal!".format(collection,
                    #                                                                               partition))
            else:
                self.logger.info(
                    'The collection or partition exist, skip this collection: {}/partition: {} '.format(collection,
                                                                                                       partition))
            del vectors
            del vector
            return return_ids
        except Exception as e:
            self.logger.error(e)
            sys.exit(2)
