package querynode

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

//----------------------------------------------------------------------------------------------------- collection
func TestCollectionReplica_getCollectionNum(t *testing.T) {
	node := newQueryNodeMock()
	initTestMeta(t, node, 0, 0)
	assert.Equal(t, node.replica.getCollectionNum(), 1)
	node.Stop()
}

func TestCollectionReplica_addCollection(t *testing.T) {
	node := newQueryNodeMock()
	initTestMeta(t, node, 0, 0)
	node.Stop()
}

func TestCollectionReplica_removeCollection(t *testing.T) {
	node := newQueryNodeMock()
	initTestMeta(t, node, 0, 0)
	assert.Equal(t, node.replica.getCollectionNum(), 1)

	err := node.replica.removeCollection(0)
	assert.NoError(t, err)
	assert.Equal(t, node.replica.getCollectionNum(), 0)
	node.Stop()
}

func TestCollectionReplica_getCollectionByID(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)
	targetCollection, err := node.replica.getCollectionByID(collectionID)
	assert.NoError(t, err)
	assert.NotNil(t, targetCollection)
	assert.Equal(t, targetCollection.ID(), collectionID)
	node.Stop()
}

func TestCollectionReplica_hasCollection(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	hasCollection := node.replica.hasCollection(collectionID)
	assert.Equal(t, hasCollection, true)
	hasCollection = node.replica.hasCollection(UniqueID(1))
	assert.Equal(t, hasCollection, false)

	node.Stop()
}

//----------------------------------------------------------------------------------------------------- partition
func TestCollectionReplica_getPartitionNum(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	partitionIDs := []UniqueID{1, 2, 3}
	for _, id := range partitionIDs {
		err := node.replica.addPartition(collectionID, id)
		assert.NoError(t, err)
		partition, err := node.replica.getPartitionByID(collectionID, id)
		assert.NoError(t, err)
		assert.Equal(t, partition.ID(), id)
	}

	partitionNum, err := node.replica.getPartitionNum(collectionID)
	assert.NoError(t, err)
	assert.Equal(t, partitionNum, len(partitionIDs)+1) // _default
	node.Stop()
}

func TestCollectionReplica_addPartition(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	partitionIDs := []UniqueID{1, 2, 3}
	for _, id := range partitionIDs {
		err := node.replica.addPartition(collectionID, id)
		assert.NoError(t, err)
		partition, err := node.replica.getPartitionByID(collectionID, id)
		assert.NoError(t, err)
		assert.Equal(t, partition.ID(), id)
	}
	node.Stop()
}

func TestCollectionReplica_removePartition(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	partitionIDs := []UniqueID{1, 2, 3}

	for _, id := range partitionIDs {
		err := node.replica.addPartition(collectionID, id)
		assert.NoError(t, err)
		partition, err := node.replica.getPartitionByID(collectionID, id)
		assert.NoError(t, err)
		assert.Equal(t, partition.ID(), id)
		err = node.replica.removePartition(collectionID, id)
		assert.NoError(t, err)
	}
	node.Stop()
}

func TestCollectionReplica_addPartitionsByCollectionMeta(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	collectionMeta := genTestCollectionMeta(collectionID, false)
	collectionMeta.PartitionIDs = []UniqueID{0, 1, 2}

	err := node.replica.addPartitionsByCollectionMeta(collectionMeta)
	assert.NoError(t, err)
	partitionNum, err := node.replica.getPartitionNum(UniqueID(0))
	assert.NoError(t, err)
	assert.Equal(t, partitionNum, len(collectionMeta.PartitionIDs)+1)
	hasPartition := node.replica.hasPartition(UniqueID(0), UniqueID(0))
	assert.Equal(t, hasPartition, true)
	hasPartition = node.replica.hasPartition(UniqueID(0), UniqueID(1))
	assert.Equal(t, hasPartition, true)
	hasPartition = node.replica.hasPartition(UniqueID(0), UniqueID(2))
	assert.Equal(t, hasPartition, true)

	node.Stop()
}

func TestCollectionReplica_removePartitionsByCollectionMeta(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	collectionMeta := genTestCollectionMeta(collectionID, false)
	collectionMeta.PartitionIDs = []UniqueID{0}

	err := node.replica.addPartitionsByCollectionMeta(collectionMeta)
	assert.NoError(t, err)
	partitionNum, err := node.replica.getPartitionNum(UniqueID(0))
	assert.NoError(t, err)
	assert.Equal(t, partitionNum, len(collectionMeta.PartitionIDs)+1)

	hasPartition := node.replica.hasPartition(UniqueID(0), UniqueID(0))
	assert.Equal(t, hasPartition, true)
	hasPartition = node.replica.hasPartition(UniqueID(0), UniqueID(1))
	assert.Equal(t, hasPartition, false)
	hasPartition = node.replica.hasPartition(UniqueID(0), UniqueID(2))
	assert.Equal(t, hasPartition, false)

	node.Stop()
}

func TestCollectionReplica_getPartitionByTag(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	collectionMeta := genTestCollectionMeta(collectionID, false)

	for _, id := range collectionMeta.PartitionIDs {
		err := node.replica.addPartition(collectionID, id)
		assert.NoError(t, err)
		partition, err := node.replica.getPartitionByID(collectionID, id)
		assert.NoError(t, err)
		assert.Equal(t, partition.ID(), id)
		assert.NotNil(t, partition)
	}
	node.Stop()
}

func TestCollectionReplica_hasPartition(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	collectionMeta := genTestCollectionMeta(collectionID, false)
	err := node.replica.addPartition(collectionID, collectionMeta.PartitionIDs[0])
	assert.NoError(t, err)
	hasPartition := node.replica.hasPartition(collectionID, defaultPartitionID)
	assert.Equal(t, hasPartition, true)
	hasPartition = node.replica.hasPartition(collectionID, defaultPartitionID+1)
	assert.Equal(t, hasPartition, false)
	node.Stop()
}

//----------------------------------------------------------------------------------------------------- segment
func TestCollectionReplica_addSegment(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	const segmentNum = 3
	for i := 0; i < segmentNum; i++ {
		err := node.replica.addSegment(UniqueID(i), defaultPartitionID, collectionID, segTypeGrowing)
		assert.NoError(t, err)
		targetSeg, err := node.replica.getSegmentByID(UniqueID(i))
		assert.NoError(t, err)
		assert.Equal(t, targetSeg.segmentID, UniqueID(i))
	}

	node.Stop()
}

func TestCollectionReplica_removeSegment(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	const segmentNum = 3

	for i := 0; i < segmentNum; i++ {
		err := node.replica.addSegment(UniqueID(i), defaultPartitionID, collectionID, segTypeGrowing)
		assert.NoError(t, err)
		targetSeg, err := node.replica.getSegmentByID(UniqueID(i))
		assert.NoError(t, err)
		assert.Equal(t, targetSeg.segmentID, UniqueID(i))
		err = node.replica.removeSegment(UniqueID(i))
		assert.NoError(t, err)
	}

	node.Stop()
}

func TestCollectionReplica_getSegmentByID(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	const segmentNum = 3

	for i := 0; i < segmentNum; i++ {
		err := node.replica.addSegment(UniqueID(i), defaultPartitionID, collectionID, segTypeGrowing)
		assert.NoError(t, err)
		targetSeg, err := node.replica.getSegmentByID(UniqueID(i))
		assert.NoError(t, err)
		assert.Equal(t, targetSeg.segmentID, UniqueID(i))
	}

	node.Stop()
}

func TestCollectionReplica_hasSegment(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	const segmentNum = 3

	for i := 0; i < segmentNum; i++ {
		err := node.replica.addSegment(UniqueID(i), defaultPartitionID, collectionID, segTypeGrowing)
		assert.NoError(t, err)
		targetSeg, err := node.replica.getSegmentByID(UniqueID(i))
		assert.NoError(t, err)
		assert.Equal(t, targetSeg.segmentID, UniqueID(i))
		hasSeg := node.replica.hasSegment(UniqueID(i))
		assert.Equal(t, hasSeg, true)
		hasSeg = node.replica.hasSegment(UniqueID(i + 100))
		assert.Equal(t, hasSeg, false)
	}

	node.Stop()
}

func TestCollectionReplica_freeAll(t *testing.T) {
	node := newQueryNodeMock()
	collectionID := UniqueID(0)
	initTestMeta(t, node, collectionID, 0)

	node.Stop()

}