import time
import random
import pdb
import threading
import logging
from multiprocessing import Pool, Process
import pytest
from milvus import IndexType, MetricType
from utils import *


dim = 128
index_file_size = 10
collection_id = "test_get_vector_by_id"
DELETE_TIMEOUT = 60
nprobe = 1
tag = "1970-01-01"
top_k = 1
nb = 6000
tag = "tag"

class TestGetBase:
    """
    ******************************************************************
      The following cases are used to test `get_vector_by_id` function
    ******************************************************************
    """
    def test_get_vector_A(self, connect, collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(collection, vector)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[0]) 
        assert status.OK()
        assert_equal_vector(res, vector[0])

    def test_get_vector_B(self, connect, collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        vectors = gen_vectors(nb, dim)
        status, ids = connect.add_vectors(collection, vectors)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[0])
        assert status.OK()
        assert_equal_vector(res, vectors[0])

    def test_get_vector_partition(self, connect, collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        vectors = gen_vectors(nb, dim)
        status = connect.create_partition(collection, tag)
        assert status.OK()
        status, ids = connect.add_vectors(collection, vectors, partition_tag=tag)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[0])
        assert status.OK()
        assert_equal_vector(res, vectors[0])

    def test_get_vector_multi_same_ids(self, connect, collection):
        '''
        target: test get_vector_by_id
        method: add vectors, with the same id, get vector by the given id
        expected: status ok, get one vector 
        '''
        vectors = gen_vectors(nb, dim)
        ids = [i for i in range(nb)]
        ids[1] = 0; ids[-1] = 0
        status, ids = connect.add_vectors(collection, vectors, ids=ids)
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, 0) 
        assert status.OK()
        assert_equal_vector(res, vectors[0])

    @pytest.fixture(
        scope="function",
        params=[
            1,
            10,
            100,
            1000,
            -1
        ],
    )
    def get_id(self, request):
        yield request.param

    def test_get_vector_after_delete(self, connect, collection, get_id):
        '''
        target: test get_vector_by_id
        method: add vectors, and delete, get vector by the given id
        expected: status ok, get one vector
        '''
        vectors = gen_vectors(nb, dim)
        status, ids = connect.add_vectors(collection, vectors)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        id = get_id
        status = connect.delete_by_id(collection, [ids[id]])
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[id])
        assert status.OK()
        assert not res 

    def test_get_vector_after_delete_with_partition(self, connect, collection, get_id):
        '''
        target: test get_vector_by_id
        method: add vectors into partition, and delete, get vector by the given id
        expected: status ok, get one vector
        '''
        vectors = gen_vectors(nb, dim)
        status = connect.create_partition(collection, tag)
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        id = get_id
        status = connect.delete_by_id(collection, [ids[id]])
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[id])
        assert status.OK()
        assert not res

    def test_get_vector_id_not_exised(self, connect, collection):
        '''
        target: test get vector, params vector_id not existed
        method: add vector and get 
        expected: status ok, empty result
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(collection, vector)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, 1) 
        assert status.OK()
        assert not res 

    def test_get_vector_collection_not_existed(self, connect, collection):
        '''
        target: test get vector, params collection_name not existed
        method: add vector and get
        expected: status not ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(collection, vector)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        collection_new = gen_unique_str()
        status, res = connect.get_vector_by_id(collection_new, 1) 
        assert not status.OK()


class TestGetIndexedVectors:
    """
    ******************************************************************
      The following cases are used to test `get_vector_by_id` function
    ******************************************************************
    """
    @pytest.fixture(
        scope="function",
        params=gen_simple_index()
    )
    def get_simple_index(self, request, connect):
        if str(connect._cmd("mode")[1]) == "CPU":
            if request.param["index_type"] not in [IndexType.IVF_SQ8, IndexType.IVFLAT, IndexType.FLAT]:
                logging.getLogger().info(request.param["index_type"])
                pytest.skip("Only support index_type: flat/ivf_flat/ivf_sq8")
        else:
            pytest.skip("Only support CPU mode")
        return request.param

    @pytest.fixture(
        scope="function",
        params=[
            1,
            10,
            100,
            1000,
            -1
        ],
    )
    def get_id(self, request):
        yield request.param

    def test_get_vectors_after_index_created(self, connect, collection, get_simple_index, get_id):
        '''
        target: test get vector after index created
        method: add vector, create index and get vector
        expected: status ok
        '''
        index_param = get_simple_index["index_param"]
        index_type = get_simple_index["index_type"]
        vectors = gen_vector(nb, dim)
        status, ids = connect.add_vectors(collection, vectors)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status = connect.create_index(collection, index_type, index_param)
        assert status.OK()
        id = get_id
        status, res = connect.get_vector_by_id(collection, ids[id])
        assert status.OK()
        logging.getLogger().info(res)
        assert status.OK()
        assert_equal_vector(res, vectors[id])

    def test_get_vector_after_delete(self, connect, collection, get_simple_index, get_id):
        '''
        target: test get_vector_by_id
        method: add vectors, and delete, get vector by the given id
        expected: status ok, get one vector
        '''
        index_param = get_simple_index["index_param"]
        index_type = get_simple_index["index_type"]
        vectors = gen_vectors(nb, dim)
        status, ids = connect.add_vectors(collection, vectors)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status = connect.create_index(collection, index_type, index_param)
        assert status.OK()
        id = get_id
        status = connect.delete_by_id(collection, [ids[id]])
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(collection, ids[id])
        assert status.OK()
        assert not res

    def test_get_vector_partition(self, connect, collection, get_simple_index, get_id):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        index_param = get_simple_index["index_param"]
        index_type = get_simple_index["index_type"]
        vectors = gen_vectors(nb, dim)
        status = connect.create_partition(collection, tag)
        ids = [i for i in range(nb)] 
        status, ids = connect.add_vectors(collection, vectors, ids, partition_tag=tag)
        assert status.OK()
        status = connect.flush([collection])
        assert status.OK()
        status = connect.create_index(collection, index_type, index_param)
        assert status.OK()
        id = get_id
        status, res = connect.get_vector_by_id(collection, ids[id])
        assert status.OK()
        assert_equal_vector(res, vectors[id])


class TestGetBinary:
    """
    ******************************************************************
      The following cases are used to test `get_vector_by_id` function
    ******************************************************************
    """
    def test_get_vector_A(self, connect, jac_collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        tmp, vector = gen_binary_vectors(1, dim)
        status, ids = connect.add_vectors(jac_collection, vector)
        assert status.OK()
        status = connect.flush([jac_collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(jac_collection, ids[0]) 
        assert status.OK()
        assert res == vector[0]

    def test_get_vector_B(self, connect, jac_collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        tmp, vectors = gen_binary_vectors(nb, dim)
        status, ids = connect.add_vectors(jac_collection, vectors)
        assert status.OK()
        status = connect.flush([jac_collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(jac_collection, ids[0])
        assert status.OK()
        assert res == vectors[0]

    def test_get_vector_multi_same_ids(self, connect, jac_collection):
        '''
        target: test get_vector_by_id
        method: add vectors, with the same id, get vector by the given id
        expected: status ok, get one vector 
        '''
        tmp, vectors = gen_binary_vectors(nb, dim)
        ids = [i for i in range(nb)]
        ids[0] = 0; ids[-1] = 0
        status, ids = connect.add_vectors(jac_collection, vectors, ids=ids)
        status = connect.flush([jac_collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(jac_collection, 0) 
        assert status.OK()
        assert res == vectors[0] 

    def test_get_vector_id_not_exised(self, connect, jac_collection):
        '''
        target: test get vector, params vector_id not existed
        method: add vector and get 
        expected: status ok, empty result
        '''
        tmp, vector = gen_binary_vectors(1, dim)
        status, ids = connect.add_vectors(jac_collection, vector)
        assert status.OK()
        status = connect.flush([jac_collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(jac_collection, 1) 
        assert status.OK()
        assert not res 

    def test_get_vector_collection_not_existed(self, connect, jac_collection):
        '''
        target: test get vector, params collection_name not existed
        method: add vector and get
        expected: status not ok
        '''
        tmp, vector = gen_binary_vectors(1, dim)
        status, ids = connect.add_vectors(jac_collection, vector)
        assert status.OK()
        status = connect.flush([jac_collection])
        assert status.OK()
        collection_new = gen_unique_str()
        status, res = connect.get_vector_by_id(collection_new, 1) 
        assert not status.OK()

    def test_get_vector_partition(self, connect, jac_collection):
        '''
        target: test get_vector_by_id
        method: add vector, and get
        expected: status ok, vector returned
        '''
        tmp, vectors = gen_binary_vectors(nb, dim)
        status = connect.create_partition(jac_collection, tag)
        status, ids = connect.add_vectors(jac_collection, vectors, partition_tag=tag)
        assert status.OK()
        status = connect.flush([jac_collection])
        assert status.OK()
        status, res = connect.get_vector_by_id(jac_collection, ids[0])
        assert status.OK()
        assert res == vectors[0]


class TestGetVectorIdIngalid(object):
    single_vector = gen_single_vector(dim)

    """
    Test adding vectors with invalid vectors
    """
    @pytest.fixture(
        scope="function",
        params=gen_invalid_vector_ids()
    )
    def gen_invalid_id(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_get_vector_id_invalid(self, connect, collection, gen_invalid_id):
        invalid_id = gen_invalid_id
        with pytest.raises(Exception) as e:
            status = connect.get_vector_by_id(collection, invalid_id)


class TestCollectionNameInvalid(object):
    """
    Test adding vectors with invalid collection names
    """
    @pytest.fixture(
        scope="function",
        params=gen_invalid_collection_names()
    )
    def get_collection_name(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_get_vectors_with_invalid_collection_name(self, connect, get_collection_name):
        collection_name = get_collection_name
        vectors = gen_vectors(1, dim)
        status, result = connect.get_vector_by_id(collection_name, 1)
        assert not status.OK()
