package datacoord

import (
	"context"
	"testing"
	"time"

	"github.com/bytedance/mockey"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
	"github.com/milvus-io/milvus/internal/datacoord/broker"
	"github.com/milvus-io/milvus/internal/metastore"
	"github.com/milvus-io/milvus/internal/metastore/model"
	mocks2 "github.com/milvus-io/milvus/internal/mocks"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/internal/storagev2/packed"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/objectstorage"
	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/workerpb"
	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
)

func TestGetQueryVChanPositionsRetrieveM2N(t *testing.T) {
	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()

	channel := "ch1"
	svr.meta.AddCollection(&collectionInfo{
		ID:         1,
		Partitions: []int64{0},
		Schema:     schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  channel,
				Data: []byte{8, 9, 10},
			},
		},
	})
	err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
		CollectionID: 1,
		FieldID:      2,
		IndexID:      1,
	})
	require.NoError(t, err)

	segArgs := []struct {
		segID   int64
		state   commonpb.SegmentState
		level   datapb.SegmentLevel
		indexed bool
	}{
		{100, commonpb.SegmentState_Growing, datapb.SegmentLevel_L1, false},
		{200, commonpb.SegmentState_Flushing, datapb.SegmentLevel_L1, false},
		{300, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, false},
		{400, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
		{401, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
		{402, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
		{403, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, true},
		{404, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, false},
		// (400[indexed], 401[indexed], 402(indexed) -> 403(indexed), 404(no index))
		{500, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L2, true},
		{600, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L2, false},
		{700, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L0, false},
		{800, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, false},
	}

	compactFroms := []int64{400, 401, 402}
	compactTos := []int64{403, 404}

	for _, arg := range segArgs {
		seg := NewSegmentInfo(&datapb.SegmentInfo{
			ID:            arg.segID,
			CollectionID:  1,
			InsertChannel: channel,
			State:         arg.state,
			Level:         arg.level,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: channel,
				MsgID:       []byte{1, 2, 3},
			},
			NumOfRows: 2048,
		})

		if lo.Contains(compactTos, arg.segID) {
			seg.CompactionFrom = compactFroms
		}
		err := svr.meta.AddSegment(context.TODO(), seg)
		require.NoError(t, err)

		if arg.indexed {
			err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
				SegmentID: arg.segID,
				BuildID:   arg.segID,
				IndexID:   1,
			})
			assert.NoError(t, err)
			err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
				BuildID: arg.segID,
				State:   commonpb.IndexState_Finished,
			})
			assert.NoError(t, err)
		}
	}

	info := svr.handler.GetQueryVChanPositions(&channelMeta{Name: channel, CollectionID: 1}, -1)

	totalSegs := len(info.GetLevelZeroSegmentIds()) +
		len(info.GetUnflushedSegmentIds()) +
		len(info.GetFlushedSegmentIds()) +
		len(info.GetDroppedSegmentIds())
	assert.EqualValues(t, 1, info.CollectionID)
	assert.EqualValues(t, len(segArgs)-2, totalSegs)
	assert.ElementsMatch(t, []int64{700}, info.GetLevelZeroSegmentIds())
	assert.ElementsMatch(t, []int64{100}, info.GetUnflushedSegmentIds())
	assert.ElementsMatch(t, []int64{200, 300, 400, 401, 402, 500, 600}, info.GetFlushedSegmentIds())
	assert.ElementsMatch(t, []int64{800}, info.GetDroppedSegmentIds())

	assert.Empty(t, info.GetUnflushedSegments())
	assert.Empty(t, info.GetFlushedSegments())
	assert.Empty(t, info.GetDroppedSegments())
	assert.Empty(t, info.GetIndexedSegments())
	assert.Empty(t, info.GetIndexedSegmentIds())
}

func TestGetQueryVChanPositions(t *testing.T) {
	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()
	svr.meta.AddCollection(&collectionInfo{
		ID:         0,
		Partitions: []int64{0, 1},
		Schema:     schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch1",
				Data: []byte{8, 9, 10},
			},
		},
	})
	svr.meta.AddCollection(&collectionInfo{
		ID:     1,
		Schema: schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch0",
				Data: []byte{8, 9, 10},
			},
		},
	})

	err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
		CollectionID: 0,
		FieldID:      2,
		IndexID:      1,
	})
	assert.NoError(t, err)

	s1 := &datapb.SegmentInfo{
		ID:            1,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   0,
		},
		NumOfRows: 2048,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s1))
	assert.NoError(t, err)
	err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
		SegmentID: 1,
		BuildID:   1,
		IndexID:   1,
	})
	assert.NoError(t, err)
	err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
		BuildID: 1,
		State:   commonpb.IndexState_Finished,
	})
	assert.NoError(t, err)
	s2 := &datapb.SegmentInfo{
		ID:            2,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Growing,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
			MsgGroup:    "",
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s2))
	assert.NoError(t, err)
	s3 := &datapb.SegmentInfo{
		ID:            3,
		CollectionID:  0,
		PartitionID:   1,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Growing,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
			MsgGroup:    "",
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{11, 12, 13},
			MsgGroup:    "",
			Timestamp:   2,
		},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s3))
	assert.NoError(t, err)

	s4 := &datapb.SegmentInfo{
		ID:            4,
		CollectionID:  0,
		PartitionID:   common.AllPartitionsID,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
			MsgGroup:    "",
			Timestamp:   1,
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{11, 12, 13},
			MsgGroup:    "",
			Timestamp:   2,
		},
		Level: datapb.SegmentLevel_L0,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s4))
	assert.NoError(t, err)

	t.Run("get unexisted channel", func(t *testing.T) {
		vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "chx1", CollectionID: 0})
		assert.Empty(t, vchan.UnflushedSegmentIds)
		assert.Empty(t, vchan.FlushedSegmentIds)
	})

	t.Run("empty collection", func(t *testing.T) {
		infos := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch0_suffix", CollectionID: 1})
		assert.EqualValues(t, 1, infos.CollectionID)
		assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
		assert.EqualValues(t, 0, len(infos.UnflushedSegmentIds))
		assert.EqualValues(t, 0, len(infos.GetLevelZeroSegmentIds()))
	})

	t.Run("filter partition", func(t *testing.T) {
		infos := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, 1)
		assert.EqualValues(t, 0, infos.CollectionID)
		assert.EqualValues(t, 1, len(infos.GetLevelZeroSegmentIds()))
		assert.EqualValues(t, uint64(1), infos.GetDeleteCheckpoint().GetTimestamp())
	})

	t.Run("empty collection with passed positions", func(t *testing.T) {
		vchannel := "ch_no_segment_1"
		pchannel := funcutil.ToPhysicalChannel(vchannel)
		infos := svr.handler.GetQueryVChanPositions(&channelMeta{
			Name:         vchannel,
			CollectionID: 0,
			StartPosition: &msgpb.MsgPosition{
				ChannelName: pchannel,
				MsgID:       []byte{14, 15, 16},
			},
		})
		assert.EqualValues(t, 0, infos.CollectionID)
		assert.EqualValues(t, vchannel, infos.ChannelName)
		assert.EqualValues(t, 0, len(infos.GetLevelZeroSegmentIds()))
	})
}

func TestGetQueryVChanPositions_PartitionStats(t *testing.T) {
	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()
	collectionID := int64(0)
	partitionID := int64(1)
	vchannel := "test_vchannel"
	version := int64(100)
	svr.meta.AddCollection(&collectionInfo{
		ID:     collectionID,
		Schema: schema,
	})
	svr.meta.partitionStatsMeta.partitionStatsInfos = map[string]map[int64]*partitionStatsInfo{
		vchannel: {
			partitionID: {
				currentVersion: version,
				infos: map[int64]*datapb.PartitionStatsInfo{
					version: {Version: version},
				},
			},
			partitionID + 1: {
				currentVersion: version + 1,
				infos: map[int64]*datapb.PartitionStatsInfo{
					version + 1: {Version: version + 1},
				},
			},
		},
	}
	vChannelInfo := svr.handler.GetQueryVChanPositions(&channelMeta{Name: vchannel, CollectionID: collectionID}, partitionID)
	statsVersions := vChannelInfo.GetPartitionStatsVersions()
	assert.Equal(t, 1, len(statsVersions))
	assert.Equal(t, int64(100), statsVersions[partitionID])

	vChannelInfo2 := svr.handler.GetQueryVChanPositions(&channelMeta{Name: vchannel, CollectionID: collectionID})
	statsVersions2 := vChannelInfo2.GetPartitionStatsVersions()
	assert.Equal(t, 2, len(statsVersions2))
}

func TestGetQueryVChanPositions_Retrieve_unIndexed(t *testing.T) {
	t.Run("ab GC-ed, cde unIndexed", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)
		schema := newTestSchema()
		svr.meta.AddCollection(&collectionInfo{
			ID:     0,
			Schema: schema,
		})
		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
			CollectionID: 0,
			FieldID:      2,
			IndexID:      1,
		})
		assert.NoError(t, err)
		c := &datapb.SegmentInfo{
			ID:            1,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
		assert.NoError(t, err)
		d := &datapb.SegmentInfo{
			ID:            2,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
		assert.NoError(t, err)
		e := &datapb.SegmentInfo{
			ID:            3,
			CollectionID:  0,
			PartitionID:   1,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{1, 2}, // c, d
			NumOfRows:      2048,
		}

		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
		assert.NoError(t, err)
		// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
		// assert.EqualValues(t, 2, len(vchan.FlushedSegmentIds))
		// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
		// assert.ElementsMatch(t, []int64{c.GetID(), d.GetID()}, vchan.FlushedSegmentIds) // expected c, d
	})

	t.Run("a GC-ed, bcde unIndexed", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)
		schema := newTestSchema()
		svr.meta.AddCollection(&collectionInfo{
			ID:     0,
			Schema: schema,
		})
		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
			CollectionID: 0,
			FieldID:      2,
			IndexID:      1,
		})
		assert.NoError(t, err)
		a := &datapb.SegmentInfo{
			ID:            99,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(a))
		assert.NoError(t, err)

		c := &datapb.SegmentInfo{
			ID:            1,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
		assert.NoError(t, err)
		d := &datapb.SegmentInfo{
			ID:            2,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
		assert.NoError(t, err)
		e := &datapb.SegmentInfo{
			ID:            3,
			CollectionID:  0,
			PartitionID:   1,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{1, 2}, // c, d
			NumOfRows:      2048,
		}

		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
		assert.NoError(t, err)
		// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
		// assert.EqualValues(t, 2, len(vchan.FlushedSegmentIds))
		// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
		// assert.ElementsMatch(t, []int64{c.GetID(), d.GetID()}, vchan.FlushedSegmentIds) // expected c, d
	})

	t.Run("ab GC-ed, c unIndexed, de indexed", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)
		schema := newTestSchema()
		svr.meta.AddCollection(&collectionInfo{
			ID:     0,
			Schema: schema,
		})
		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
			CollectionID: 0,
			FieldID:      2,
			IndexID:      1,
		})
		assert.NoError(t, err)
		c := &datapb.SegmentInfo{
			ID:            1,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
		assert.NoError(t, err)
		d := &datapb.SegmentInfo{
			ID:            2,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
		assert.NoError(t, err)
		err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
			SegmentID: 2,
			BuildID:   1,
			IndexID:   1,
		})
		assert.NoError(t, err)
		err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
			BuildID: 1,
			State:   commonpb.IndexState_Finished,
		})
		assert.NoError(t, err)
		e := &datapb.SegmentInfo{
			ID:            3,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			CompactionFrom: []int64{1, 2}, // c, d
			NumOfRows:      2048,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
		assert.NoError(t, err)
		err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
			SegmentID: 3,
			BuildID:   2,
			IndexID:   1,
		})
		assert.NoError(t, err)
		err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
			BuildID: 2,
			State:   commonpb.IndexState_Finished,
		})
		assert.NoError(t, err)

		// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
		// assert.EqualValues(t, 1, len(vchan.FlushedSegmentIds))
		// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
		// assert.ElementsMatch(t, []int64{e.GetID()}, vchan.FlushedSegmentIds) // expected e
	})

	t.Run("complex derivation", func(t *testing.T) {
		// numbers indicate segmentID, letters indicate segment information
		// i: indexed,  u: unindexed,  g: gced
		//        1i,  2i,  3g     4i,  5i,  6i
		//        |    |    |      |    |    |
		//        \    |    /      \    |    /
		//         \   |   /        \   |   /
		//    7u,  [8i,9i,10i]      [11u, 12i]
		//    |     |  |  |          |     |
		//    \     |  /  \          /     |
		//     \    | /    \        /      |
		//       [13u]      [14i, 15u]    12i                       [19u](unsorted)
		//        |         |     |        |                              |
		//        \         /     \        /                              |
		//         \       /       \      /                               |
		//           [16u]          [17u]       [18u](unsorted)     [20u](sorted)      [21i](unsorted)
		// all leaf nodes are [1,2,3,4,5,6,7], but because segment3 has been gced, the leaf node becomes [7,8,9,10,4,5,6]
		// should be returned: flushed: [7, 8, 9, 10, 4, 5, 6, 20, 21], growing: [18]
		svr := newTestServer(t)
		defer closeTestServer(t, svr)
		schema := newTestSchema()
		svr.meta.AddCollection(&collectionInfo{
			ID:         0,
			Partitions: []int64{0},
			Schema:     schema,
		})
		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
			CollectionID: 0,
			FieldID:      2,
			IndexID:      1,
		})
		assert.NoError(t, err)
		seg1 := &datapb.SegmentInfo{
			ID:            1,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
		assert.NoError(t, err)
		seg2 := &datapb.SegmentInfo{
			ID:            2,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
		assert.NoError(t, err)
		// seg3 was GCed
		seg4 := &datapb.SegmentInfo{
			ID:            4,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
		assert.NoError(t, err)
		seg5 := &datapb.SegmentInfo{
			ID:            5,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
		assert.NoError(t, err)
		seg6 := &datapb.SegmentInfo{
			ID:            6,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
		assert.NoError(t, err)
		seg7 := &datapb.SegmentInfo{
			ID:            7,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 2048,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
		assert.NoError(t, err)
		seg8 := &datapb.SegmentInfo{
			ID:            8,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{1, 2, 3},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
		assert.NoError(t, err)
		seg9 := &datapb.SegmentInfo{
			ID:            9,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{1, 2, 3},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
		assert.NoError(t, err)
		seg10 := &datapb.SegmentInfo{
			ID:            10,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{1, 2, 3},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg10))
		assert.NoError(t, err)
		seg11 := &datapb.SegmentInfo{
			ID:            11,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2048,
			CompactionFrom:      []int64{4, 5, 6},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg11))
		assert.NoError(t, err)
		seg12 := &datapb.SegmentInfo{
			ID:            12,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{4, 5, 6},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg12))
		assert.NoError(t, err)
		seg13 := &datapb.SegmentInfo{
			ID:            13,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2047,
			CompactionFrom:      []int64{7, 8, 9},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg13))
		assert.NoError(t, err)
		seg14 := &datapb.SegmentInfo{
			ID:            14,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{10, 11},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg14))
		assert.NoError(t, err)
		seg15 := &datapb.SegmentInfo{
			ID:            15,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2048,
			CompactionFrom:      []int64{10, 11},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg15))
		assert.NoError(t, err)
		seg16 := &datapb.SegmentInfo{
			ID:            16,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2048,
			CompactionFrom:      []int64{13, 14},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg16))
		assert.NoError(t, err)
		seg17 := &datapb.SegmentInfo{
			ID:            17,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2048,
			CompactionFrom:      []int64{12, 15},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg17))
		assert.NoError(t, err)
		seg18 := &datapb.SegmentInfo{
			ID:            18,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			Level:         datapb.SegmentLevel_L1,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:      2048,
			CompactionFrom: []int64{},
			IsInvisible:    true,
			IsSorted:       false,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg18))
		assert.NoError(t, err)
		seg19 := &datapb.SegmentInfo{
			ID:            19,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			Level:         datapb.SegmentLevel_L1,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:      2048,
			CompactionFrom: []int64{},
			IsInvisible:    true,
			IsSorted:       false,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg19))
		assert.NoError(t, err)
		seg20 := &datapb.SegmentInfo{
			ID:            20,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			Level:         datapb.SegmentLevel_L1,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           2048,
			CompactionFrom:      []int64{19},
			CreatedByCompaction: true,
			IsInvisible:         false,
			IsSorted:            true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg20))
		assert.NoError(t, err)
		seg21 := &datapb.SegmentInfo{
			ID:            21,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			Level:         datapb.SegmentLevel_L1,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:      100,
			CompactionFrom: []int64{},
			IsInvisible:    false,
			IsSorted:       false,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg21))
		assert.NoError(t, err)

		vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
		assert.ElementsMatch(t, []int64{7, 8, 9, 10, 4, 5, 6, 20, 21}, vchan.FlushedSegmentIds)
		assert.ElementsMatch(t, []int64{18}, vchan.UnflushedSegmentIds)
		assert.ElementsMatch(t, []int64{1, 2, 19}, vchan.DroppedSegmentIds)
	})

	t.Run("compaction iterate", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)
		schema := newTestSchema()
		svr.meta.AddCollection(&collectionInfo{
			ID:         0,
			Partitions: []int64{0},
			Schema:     schema,
		})
		err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
			CollectionID: 0,
			FieldID:      2,
			IndexID:      1,
		})
		assert.NoError(t, err)
		seg1 := &datapb.SegmentInfo{
			ID:            1,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
		assert.NoError(t, err)
		seg2 := &datapb.SegmentInfo{
			ID:            2,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows: 100,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
		assert.NoError(t, err)
		seg3 := &datapb.SegmentInfo{
			ID:            3,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{1, 2},
			IsInvisible:         true,
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg3))
		assert.NoError(t, err)
		seg4 := &datapb.SegmentInfo{
			ID:            4,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{1, 2},
			IsInvisible:         true,
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
		assert.NoError(t, err)
		seg5 := &datapb.SegmentInfo{
			ID:            5,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{3},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
		assert.NoError(t, err)
		seg6 := &datapb.SegmentInfo{
			ID:            6,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{4},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
		assert.NoError(t, err)
		seg7 := &datapb.SegmentInfo{
			ID:            7,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{5, 6},
			IsInvisible:         true,
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
		assert.NoError(t, err)
		seg8 := &datapb.SegmentInfo{
			ID:            8,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Dropped,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{5, 6},
			IsInvisible:         true,
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
		assert.NoError(t, err)
		seg9 := &datapb.SegmentInfo{
			ID:            9,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{7},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
		assert.NoError(t, err)
		seg10 := &datapb.SegmentInfo{
			ID:            10,
			CollectionID:  0,
			PartitionID:   0,
			InsertChannel: "ch1",
			State:         commonpb.SegmentState_Flushed,
			DmlPosition: &msgpb.MsgPosition{
				ChannelName: "ch1",
				MsgID:       []byte{1, 2, 3},
				MsgGroup:    "",
				Timestamp:   1,
			},
			NumOfRows:           100,
			CompactionFrom:      []int64{8},
			CreatedByCompaction: true,
		}
		err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg10))
		assert.NoError(t, err)

		vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
		assert.ElementsMatch(t, []int64{5, 6}, vchan.FlushedSegmentIds)
		assert.ElementsMatch(t, []int64{}, vchan.UnflushedSegmentIds)
		assert.ElementsMatch(t, []int64{1, 2}, vchan.DroppedSegmentIds)
	})
}

func TestGetCurrentSegmentsView(t *testing.T) {
	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()
	svr.meta.AddCollection(&collectionInfo{
		ID:         0,
		Partitions: []int64{0},
		Schema:     schema,
	})

	err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
		CollectionID: 0,
		FieldID:      2,
		IndexID:      1,
	})
	assert.NoError(t, err)
	seg1 := &datapb.SegmentInfo{
		ID:            1,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Dropped,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows: 100,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
	assert.NoError(t, err)
	seg2 := &datapb.SegmentInfo{
		ID:            2,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows:      100,
		CompactionFrom: []int64{1},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
	assert.NoError(t, err)
	seg3 := &datapb.SegmentInfo{
		ID:            3,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows: 100,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg3))
	assert.NoError(t, err)
	seg4 := &datapb.SegmentInfo{
		ID:            4,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows: 100,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
	assert.NoError(t, err)
	seg5 := &datapb.SegmentInfo{
		ID:            5,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows:      100,
		CompactionFrom: []int64{3, 4},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
	assert.NoError(t, err)
	seg6 := &datapb.SegmentInfo{
		ID:            6,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows:      100,
		CompactionFrom: []int64{3, 4},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
	assert.NoError(t, err)
	seg7 := &datapb.SegmentInfo{
		ID:            7,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		Level:         datapb.SegmentLevel_L0,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
	assert.NoError(t, err)
	seg8 := &datapb.SegmentInfo{
		ID:            8,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Growing,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows: 100,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
	assert.NoError(t, err)
	seg9 := &datapb.SegmentInfo{
		ID:            9,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Importing,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			MsgGroup:    "",
			Timestamp:   1,
		},
		NumOfRows: 100,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
	assert.NoError(t, err)

	view := svr.handler.GetCurrentSegmentsView(context.Background(), &channelMeta{Name: "ch1", CollectionID: 0})
	assert.ElementsMatch(t, []int64{2, 3, 4}, view.FlushedSegmentIDs)
	assert.ElementsMatch(t, []int64{8}, view.GrowingSegmentIDs)
	assert.ElementsMatch(t, []int64{1}, view.DroppedSegmentIDs)
	assert.ElementsMatch(t, []int64{7}, view.L0SegmentIDs)
	assert.ElementsMatch(t, []int64{9}, view.ImportingSegmentIDs)
}

func TestShouldDropChannel(t *testing.T) {
	type myRootCoord struct {
		mocks2.MockRootCoordClient
	}
	myRoot := &myRootCoord{}
	myRoot.EXPECT().AllocTimestamp(mock.Anything, mock.Anything).Return(&rootcoordpb.AllocTimestampResponse{
		Status:    merr.Success(),
		Timestamp: tsoutil.ComposeTSByTime(time.Now()),
		Count:     1,
	}, nil)

	myRoot.EXPECT().AllocID(mock.Anything, mock.Anything).Return(&rootcoordpb.AllocIDResponse{
		Status: merr.Success(),
		ID:     int64(tsoutil.ComposeTSByTime(time.Now())),
		Count:  1,
	}, nil)

	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()
	svr.meta.AddCollection(&collectionInfo{
		ID:     0,
		Schema: schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch1",
				Data: []byte{8, 9, 10},
			},
		},
	})
	svr.meta.AddCollection(&collectionInfo{
		ID:     1,
		Schema: schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch0",
				Data: []byte{8, 9, 10},
			},
		},
	})

	t.Run("channel name not in kv ", func(t *testing.T) {
		assert.False(t, svr.handler.CheckShouldDropChannel("ch99"))
	})

	t.Run("channel in remove flag", func(t *testing.T) {
		err := svr.meta.catalog.Update(context.TODO(), metastore.MarkChannelDropped("ch1"))
		require.NoError(t, err)
		assert.True(t, svr.handler.CheckShouldDropChannel("ch1"))
	})
}

func TestGetDataVChanPositions(t *testing.T) {
	svr := newTestServer(t)
	defer closeTestServer(t, svr)
	schema := newTestSchema()
	svr.meta.AddCollection(&collectionInfo{
		ID:     0,
		Schema: schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch1",
				Data: []byte{8, 9, 10},
			},
		},
	})
	svr.meta.AddCollection(&collectionInfo{
		ID:     1,
		Schema: schema,
		StartPositions: []*commonpb.KeyDataPair{
			{
				Key:  "ch0",
				Data: []byte{8, 9, 10},
			},
		},
	})

	s1 := &datapb.SegmentInfo{
		ID:            1,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
		},
	}
	err := svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s1))
	require.Nil(t, err)
	s2 := &datapb.SegmentInfo{
		ID:            2,
		CollectionID:  0,
		PartitionID:   0,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Growing,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{1, 2, 3},
			Timestamp:   1,
		},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s2))
	require.Nil(t, err)
	s3 := &datapb.SegmentInfo{
		ID:            3,
		CollectionID:  0,
		PartitionID:   1,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Growing,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{11, 12, 13},
			Timestamp:   2,
		},
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s3))
	require.Nil(t, err)
	l0Segment := &datapb.SegmentInfo{
		ID:            4,
		CollectionID:  0,
		PartitionID:   1,
		InsertChannel: "ch1",
		State:         commonpb.SegmentState_Flushed,
		StartPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{8, 9, 10},
		},
		DmlPosition: &msgpb.MsgPosition{
			ChannelName: "ch1",
			MsgID:       []byte{11, 12, 13},
			Timestamp:   2,
		},
		Level: datapb.SegmentLevel_L0,
	}
	err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(l0Segment))
	require.Nil(t, err)

	t.Run("get unexisted channel", func(t *testing.T) {
		vchan := svr.handler.GetDataVChanPositions(&channelMeta{Name: "chx1", CollectionID: 0}, allPartitionID)
		assert.Empty(t, vchan.UnflushedSegmentIds)
		assert.Empty(t, vchan.FlushedSegmentIds)
	})

	t.Run("get existed channel", func(t *testing.T) {
		vchan := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, allPartitionID)
		assert.EqualValues(t, 1, len(vchan.FlushedSegmentIds))
		assert.EqualValues(t, 1, vchan.FlushedSegmentIds[0])
		assert.EqualValues(t, 2, len(vchan.UnflushedSegmentIds))
		assert.ElementsMatch(t, []int64{s2.ID, s3.ID}, vchan.UnflushedSegmentIds)
		assert.EqualValues(t, 1, len(vchan.LevelZeroSegmentIds))
		assert.ElementsMatch(t, []int64{l0Segment.ID}, vchan.LevelZeroSegmentIds)
	})

	t.Run("empty collection", func(t *testing.T) {
		infos := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch0_suffix", CollectionID: 1}, allPartitionID)
		assert.EqualValues(t, 1, infos.CollectionID)
		assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
		assert.EqualValues(t, 0, len(infos.UnflushedSegmentIds))
	})

	t.Run("filter partition", func(t *testing.T) {
		infos := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, 1)
		assert.EqualValues(t, 0, infos.CollectionID)
		assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
		assert.EqualValues(t, 1, len(infos.UnflushedSegmentIds))
	})

	t.Run("empty collection with passed positions", func(t *testing.T) {
		vchannel := "ch_no_segment_1"
		pchannel := funcutil.ToPhysicalChannel(vchannel)
		infos := svr.handler.GetDataVChanPositions(&channelMeta{
			Name:         vchannel,
			CollectionID: 0,
			StartPosition: &msgpb.MsgPosition{
				ChannelName: pchannel,
				MsgID:       []byte{14, 15, 16},
			},
		}, allPartitionID)
		assert.EqualValues(t, 0, infos.CollectionID)
		assert.EqualValues(t, vchannel, infos.ChannelName)
	})
}

func TestGetSnapshotSeekPositions_ReturnsSortedPositionsAndMinTs(t *testing.T) {
	handler := &ServerHandler{s: &Server{}}

	mockChannels := mockey.Mock((*Server).getChannelsByCollectionID).To(func(s *Server, ctx context.Context, collectionID int64) ([]RWChannel, error) {
		assert.Equal(t, UniqueID(100), collectionID)
		return []RWChannel{
			&channelMeta{Name: "ch-b", CollectionID: 100},
			&channelMeta{Name: "ch-a", CollectionID: 100},
		}, nil
	}).Build()
	defer mockChannels.UnPatch()

	chAPosition := &msgpb.MsgPosition{
		Timestamp: 100,
		MsgID:     []byte{1},
	}
	chBPosition := &msgpb.MsgPosition{
		ChannelName: "ch-b",
		Timestamp:   1000,
		MsgID:       []byte{2},
	}
	mockSeekPosition := mockey.Mock((*ServerHandler).GetChannelSeekPosition).To(func(h *ServerHandler, channel RWChannel, partitionIDs ...UniqueID) *msgpb.MsgPosition {
		assert.Equal(t, []UniqueID{10, 20}, partitionIDs)
		switch channel.GetName() {
		case "ch-a":
			return chAPosition
		case "ch-b":
			return chBPosition
		default:
			t.Fatalf("unexpected channel: %s", channel.GetName())
			return nil
		}
	}).Build()
	defer mockSeekPosition.UnPatch()

	positions, minTs, err := handler.GetSnapshotSeekPositions(context.Background(), 100, 10, 20)
	require.NoError(t, err)
	assert.Equal(t, uint64(100), minTs)
	require.Len(t, positions, 2)
	assert.Equal(t, "ch-a", positions[0].GetChannelName())
	assert.Equal(t, uint64(100), positions[0].GetTimestamp())
	assert.Equal(t, []byte{1}, positions[0].GetMsgID())
	assert.Equal(t, "ch-b", positions[1].GetChannelName())
	assert.Equal(t, uint64(1000), positions[1].GetTimestamp())
	assert.Equal(t, []byte{2}, positions[1].GetMsgID())

	positions[0].Timestamp = 9999

	positions, minTs, err = handler.GetSnapshotSeekPositions(context.Background(), 100, 10, 20)
	require.NoError(t, err)
	assert.Equal(t, uint64(100), minTs)
	require.Len(t, positions, 2)
	assert.Equal(t, "ch-a", positions[0].GetChannelName())
	assert.Equal(t, uint64(100), positions[0].GetTimestamp())
}

func TestGetDeltaLogFromCompactTo(t *testing.T) {
	t.Run("no compaction children", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)

		// Add a segment with no compaction children
		seg := NewSegmentInfo(&datapb.SegmentInfo{
			ID:           1001,
			CollectionID: 100,
			PartitionID:  10,
			State:        commonpb.SegmentState_Flushed,
		})
		err := svr.meta.AddSegment(context.TODO(), seg)
		require.NoError(t, err)

		// Should return empty delta logs
		deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1001)
		assert.NoError(t, err)
		assert.Empty(t, deltaLogs)
	})

	t.Run("segment not found", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)

		// Query non-existent segment
		_, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 9999)
		assert.Error(t, err)
	})

	t.Run("single level compaction", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)

		// Create parent segment
		parent := NewSegmentInfo(&datapb.SegmentInfo{
			ID:           1000,
			CollectionID: 100,
			PartitionID:  10,
			State:        commonpb.SegmentState_Dropped,
		})
		err := svr.meta.AddSegment(context.TODO(), parent)
		require.NoError(t, err)

		// Create child segment with delta logs
		child := NewSegmentInfo(&datapb.SegmentInfo{
			ID:             1001,
			CollectionID:   100,
			PartitionID:    10,
			State:          commonpb.SegmentState_Flushed,
			CompactionFrom: []int64{1000},
			Deltalogs: []*datapb.FieldBinlog{
				{
					FieldID: 0,
					Binlogs: []*datapb.Binlog{
						{LogID: 1, LogSize: 100},
					},
				},
			},
		})
		err = svr.meta.AddSegment(context.TODO(), child)
		require.NoError(t, err)

		// Should return child's delta logs
		deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1000)
		assert.NoError(t, err)
		assert.Len(t, deltaLogs, 1)
		assert.Equal(t, int64(0), deltaLogs[0].FieldID)
	})

	t.Run("multi-level compaction", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)

		// Create compaction chain: seg1 -> seg2 -> seg3
		seg1 := NewSegmentInfo(&datapb.SegmentInfo{
			ID:           1001,
			CollectionID: 100,
			PartitionID:  10,
			State:        commonpb.SegmentState_Dropped,
		})
		err := svr.meta.AddSegment(context.TODO(), seg1)
		require.NoError(t, err)

		seg2 := NewSegmentInfo(&datapb.SegmentInfo{
			ID:             1002,
			CollectionID:   100,
			PartitionID:    10,
			State:          commonpb.SegmentState_Dropped,
			CompactionFrom: []int64{1001},
			Deltalogs: []*datapb.FieldBinlog{
				{
					FieldID: 0,
					Binlogs: []*datapb.Binlog{
						{LogID: 2, LogSize: 100},
					},
				},
			},
		})
		err = svr.meta.AddSegment(context.TODO(), seg2)
		require.NoError(t, err)

		seg3 := NewSegmentInfo(&datapb.SegmentInfo{
			ID:             1003,
			CollectionID:   100,
			PartitionID:    10,
			State:          commonpb.SegmentState_Flushed,
			CompactionFrom: []int64{1002},
			Deltalogs: []*datapb.FieldBinlog{
				{
					FieldID: 0,
					Binlogs: []*datapb.Binlog{
						{LogID: 3, LogSize: 100},
					},
				},
			},
		})
		err = svr.meta.AddSegment(context.TODO(), seg3)
		require.NoError(t, err)

		// Should return all delta logs from the compaction chain
		deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1001)
		assert.NoError(t, err)
		assert.Len(t, deltaLogs, 2) // seg2 and seg3 delta logs
	})

	t.Run("multi-child compaction", func(t *testing.T) {
		svr := newTestServer(t)
		defer closeTestServer(t, svr)

		// Create parent segment
		parent := NewSegmentInfo(&datapb.SegmentInfo{
			ID:           1000,
			CollectionID: 100,
			PartitionID:  10,
			State:        commonpb.SegmentState_Dropped,
		})
		err := svr.meta.AddSegment(context.TODO(), parent)
		require.NoError(t, err)

		// Create multiple children from the same parent
		child1 := NewSegmentInfo(&datapb.SegmentInfo{
			ID:             1001,
			CollectionID:   100,
			PartitionID:    10,
			State:          commonpb.SegmentState_Flushed,
			CompactionFrom: []int64{1000},
			Deltalogs: []*datapb.FieldBinlog{
				{
					FieldID: 0,
					Binlogs: []*datapb.Binlog{
						{LogID: 1, LogSize: 100},
					},
				},
			},
		})
		err = svr.meta.AddSegment(context.TODO(), child1)
		require.NoError(t, err)

		child2 := NewSegmentInfo(&datapb.SegmentInfo{
			ID:             1002,
			CollectionID:   100,
			PartitionID:    10,
			State:          commonpb.SegmentState_Flushed,
			CompactionFrom: []int64{1000},
			Deltalogs: []*datapb.FieldBinlog{
				{
					FieldID: 0,
					Binlogs: []*datapb.Binlog{
						{LogID: 2, LogSize: 100},
					},
				},
			},
		})
		err = svr.meta.AddSegment(context.TODO(), child2)
		require.NoError(t, err)

		// Should return delta logs from both children
		deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1000)
		assert.NoError(t, err)
		assert.Len(t, deltaLogs, 2)
	})
}

func TestGenSnapshot(t *testing.T) {
	schema := newTestSchema()

	// Create minimal components for testing
	mockIndexMeta := &indexMeta{}
	mockMeta := &meta{
		indexMeta: mockIndexMeta,
	}

	// Create mock mixCoord
	mockMixCoord := mocks2.NewMixCoord(t)
	mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
		Status:              merr.Success(),
		Schema:              schema,
		ShardsNum:           2,
		NumPartitions:       1,
		ConsistencyLevel:    commonpb.ConsistencyLevel_Strong,
		Properties:          []*commonpb.KeyValuePair{},
		CollectionID:        200,
		VirtualChannelNames: []string{"dml_0_200v0", "dml_1_200v1"},
	}, nil).Maybe()

	mockMixCoord.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
		Status:         merr.Success(),
		PartitionIDs:   []int64{0, 1},
		PartitionNames: []string{"_default", "partition1"},
	}, nil).Maybe()

	// Create handler with mock broker
	mockBroker := broker.NewCoordinatorBroker(mockMixCoord)
	handler := &ServerHandler{
		s: &Server{
			broker: mockBroker,
			meta:   mockMeta,
		},
	}

	// Setup mocks for other methods
	mock1 := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
		return []*msgpb.MsgPosition{
			{ChannelName: "dml_0_200v0", Timestamp: 12345, MsgID: []byte{1}},
		}, uint64(12345), nil
	}).Build()
	defer mock1.UnPatch()

	mock4 := mockey.Mock((*indexMeta).GetIndexesForCollection).To(func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
		return []*model.Index{
			{
				IndexID:      1,
				CollectionID: 200,
				FieldID:      2,
				IndexName:    "test_index",
				CreateTime:   10000, // Earlier than snapshot timestamp
			},
		}
	}).Build()
	defer mock4.UnPatch()

	mock5 := mockey.Mock((*meta).SelectSegments).To(func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
		seg := NewSegmentInfo(&datapb.SegmentInfo{
			ID:            1001,
			CollectionID:  200,
			PartitionID:   0,
			InsertChannel: "dml_0_200v0",
			State:         commonpb.SegmentState_Flushed,
			StartPosition: &msgpb.MsgPosition{Timestamp: 10000},
			Binlogs: []*datapb.FieldBinlog{
				{
					FieldID: 1,
					Binlogs: []*datapb.Binlog{
						{LogID: 1, LogSize: 100},
					},
				},
			},
		})
		return []*SegmentInfo{seg}
	}).Build()
	defer mock5.UnPatch()

	mock6 := mockey.Mock((*meta).GetSegment).To(func(m *meta, ctx context.Context, segmentID int64) *SegmentInfo {
		if segmentID == 1001 {
			return NewSegmentInfo(&datapb.SegmentInfo{
				ID:           1001,
				CollectionID: 200,
				PartitionID:  0,
				State:        commonpb.SegmentState_Flushed,
				Binlogs: []*datapb.FieldBinlog{
					{
						FieldID: 1,
						Binlogs: []*datapb.Binlog{
							{LogID: 1, LogSize: 100},
						},
					},
				},
				Deltalogs: []*datapb.FieldBinlog{},
			})
		}
		return nil
	}).Build()
	defer mock6.UnPatch()

	mock7 := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
		return map[int64]*model.SegmentIndex{}
	}).Build()
	defer mock7.UnPatch()

	mock8 := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
		return nil, false // No compaction children
	}).Build()
	defer mock8.UnPatch()

	// Test GenSnapshot
	snapshotData, err := handler.GenSnapshot(context.Background(), 200)
	assert.NoError(t, err)
	assert.NotNil(t, snapshotData)
	assert.Equal(t, int64(200), snapshotData.SnapshotInfo.CollectionId)
	assert.Equal(t, int64(12345), snapshotData.SnapshotInfo.CreateTs)
	assert.NotNil(t, snapshotData.Collection)
	assert.Equal(t, int64(2), snapshotData.Collection.NumShards)
	assert.Equal(t, 1, len(snapshotData.Indexes))
	assert.Equal(t, 1, len(snapshotData.Segments))
	assert.Equal(t, int64(1001), snapshotData.Segments[0].SegmentId)
	// Verify VirtualChannelNames is populated from DescribeCollectionInternal response
	assert.Equal(t, []string{"dml_0_200v0", "dml_1_200v1"}, snapshotData.Collection.VirtualChannelNames)
}

func TestGenSnapshot_UsesPerChannelSeekPositions(t *testing.T) {
	schema := newTestSchema()
	segCh1 := NewSegmentInfo(&datapb.SegmentInfo{
		ID:            1001,
		CollectionID:  200,
		PartitionID:   0,
		InsertChannel: "ch-1",
		State:         commonpb.SegmentState_Flushed,
		StartPosition: &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 500},
		DmlPosition:   &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 600},
		Binlogs: []*datapb.FieldBinlog{
			{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
		},
	})
	segCh2 := NewSegmentInfo(&datapb.SegmentInfo{
		ID:            1002,
		CollectionID:  200,
		PartitionID:   0,
		InsertChannel: "ch-2",
		State:         commonpb.SegmentState_Flushed,
		StartPosition: &msgpb.MsgPosition{ChannelName: "ch-2", Timestamp: 500},
		DmlPosition:   &msgpb.MsgPosition{ChannelName: "ch-2", Timestamp: 600},
		Binlogs: []*datapb.FieldBinlog{
			{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 2, LogSize: 100}}},
		},
	})

	mockMeta := &meta{indexMeta: &indexMeta{}}
	handler := &ServerHandler{
		s: &Server{
			broker: broker.NewCoordinatorBroker(newMockMixCoord()),
			meta:   mockMeta,
		},
	}

	mockDescribe := mockey.Mock((*mockMixCoord).DescribeCollectionInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
			return &milvuspb.DescribeCollectionResponse{
				Status:              merr.Success(),
				Schema:              schema,
				ShardsNum:           2,
				NumPartitions:       1,
				ConsistencyLevel:    commonpb.ConsistencyLevel_Strong,
				CollectionID:        200,
				VirtualChannelNames: []string{"ch-1", "ch-2"},
			}, nil
		}).Build()
	defer mockDescribe.UnPatch()

	mockShowPartitions := mockey.Mock((*mockMixCoord).ShowPartitionsInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
			return &milvuspb.ShowPartitionsResponse{
				Status:         merr.Success(),
				PartitionIDs:   []int64{0},
				PartitionNames: []string{"_default"},
			}, nil
		}).Build()
	defer mockShowPartitions.UnPatch()

	mockSeekPositions := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(
		func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
			return []*msgpb.MsgPosition{
				{ChannelName: "ch-1", Timestamp: 100, MsgID: []byte{1}},
				{ChannelName: "ch-2", Timestamp: 1000, MsgID: []byte{2}},
			}, uint64(100), nil
		}).Build()
	defer mockSeekPositions.UnPatch()

	mockIndexes := mockey.Mock((*indexMeta).GetIndexesForCollection).To(
		func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
			return []*model.Index{}
		}).Build()
	defer mockIndexes.UnPatch()

	mockSelectSegments := mockey.Mock((*meta).SelectSegments).To(
		func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
			candidates := []*SegmentInfo{segCh1, segCh2}
			var result []*SegmentInfo
			for _, seg := range candidates {
				pass := true
				for _, filter := range filters {
					if !filter.Match(seg) {
						pass = false
						break
					}
				}
				if pass {
					result = append(result, seg)
				}
			}
			return result
		}).Build()
	defer mockSelectSegments.UnPatch()

	mockCompactionTo := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
		return nil, false
	}).Build()
	defer mockCompactionTo.UnPatch()

	mockSegmentIndexes := mockey.Mock((*indexMeta).getSegmentIndexes).To(
		func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
			return map[int64]*model.SegmentIndex{}
		}).Build()
	defer mockSegmentIndexes.UnPatch()

	snapshotData, err := handler.GenSnapshot(context.Background(), 200)
	require.NoError(t, err)
	require.NotNil(t, snapshotData)
	require.NotNil(t, snapshotData.SnapshotInfo)
	assert.Equal(t, int64(100), snapshotData.SnapshotInfo.CreateTs)
	require.Len(t, snapshotData.SnapshotInfo.GetChannelSeekPositions(), 2)
	assert.Equal(t, "ch-1", snapshotData.SnapshotInfo.GetChannelSeekPositions()[0].GetChannelName())
	assert.Equal(t, uint64(100), snapshotData.SnapshotInfo.GetChannelSeekPositions()[0].GetTimestamp())
	assert.Equal(t, "ch-2", snapshotData.SnapshotInfo.GetChannelSeekPositions()[1].GetChannelName())
	assert.Equal(t, uint64(1000), snapshotData.SnapshotInfo.GetChannelSeekPositions()[1].GetTimestamp())
	require.Len(t, snapshotData.Segments, 1)
	assert.Equal(t, int64(1002), snapshotData.Segments[0].GetSegmentId())
	assert.Equal(t, "ch-2", snapshotData.Segments[0].GetChannelName())
}

// TestGenSnapshot_IncludesV3ManifestOnlySegment guards the snapshot filter: a
// V3 segment that reloaded from etcd with empty Binlogs/Deltalogs (per-field
// KVs are not persisted for V3) but a non-empty ManifestPath still carries
// data and must not be dropped from the snapshot. An empty segment with no
// manifest is the negative control.
func TestGenSnapshot_IncludesV3ManifestOnlySegment(t *testing.T) {
	schema := newTestSchema()
	// V3 segment: empty binlog arrays but a live manifest path — must be kept.
	manifestSeg := NewSegmentInfo(&datapb.SegmentInfo{
		ID:             2001,
		CollectionID:   200,
		PartitionID:    0,
		InsertChannel:  "ch-1",
		State:          commonpb.SegmentState_Flushed,
		StartPosition:  &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 500},
		DmlPosition:    &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 600},
		StorageVersion: storage.StorageV3,
		ManifestPath:   packed.MarshalManifestPath("/data/segments/2001", 3),
	})
	// Empty segment with no manifest — must be dropped.
	emptySeg := NewSegmentInfo(&datapb.SegmentInfo{
		ID:            2002,
		CollectionID:  200,
		PartitionID:   0,
		InsertChannel: "ch-1",
		State:         commonpb.SegmentState_Flushed,
		StartPosition: &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 500},
		DmlPosition:   &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 600},
	})

	mockMeta := &meta{indexMeta: &indexMeta{}}
	handler := &ServerHandler{
		s: &Server{
			broker: broker.NewCoordinatorBroker(newMockMixCoord()),
			meta:   mockMeta,
		},
	}

	mockDescribe := mockey.Mock((*mockMixCoord).DescribeCollectionInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
			return &milvuspb.DescribeCollectionResponse{
				Status:              merr.Success(),
				Schema:              schema,
				ShardsNum:           1,
				NumPartitions:       1,
				ConsistencyLevel:    commonpb.ConsistencyLevel_Strong,
				CollectionID:        200,
				VirtualChannelNames: []string{"ch-1"},
			}, nil
		}).Build()
	defer mockDescribe.UnPatch()

	mockShowPartitions := mockey.Mock((*mockMixCoord).ShowPartitionsInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
			return &milvuspb.ShowPartitionsResponse{
				Status:         merr.Success(),
				PartitionIDs:   []int64{0},
				PartitionNames: []string{"_default"},
			}, nil
		}).Build()
	defer mockShowPartitions.UnPatch()

	mockSeekPositions := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(
		func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
			return []*msgpb.MsgPosition{
				{ChannelName: "ch-1", Timestamp: 1000, MsgID: []byte{1}},
			}, uint64(1000), nil
		}).Build()
	defer mockSeekPositions.UnPatch()

	mockIndexes := mockey.Mock((*indexMeta).GetIndexesForCollection).To(
		func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
			return []*model.Index{}
		}).Build()
	defer mockIndexes.UnPatch()

	mockSelectSegments := mockey.Mock((*meta).SelectSegments).To(
		func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
			candidates := []*SegmentInfo{manifestSeg, emptySeg}
			var result []*SegmentInfo
			for _, seg := range candidates {
				pass := true
				for _, filter := range filters {
					if !filter.Match(seg) {
						pass = false
						break
					}
				}
				if pass {
					result = append(result, seg)
				}
			}
			return result
		}).Build()
	defer mockSelectSegments.UnPatch()

	mockCompactionTo := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
		return nil, false
	}).Build()
	defer mockCompactionTo.UnPatch()

	mockSegmentIndexes := mockey.Mock((*indexMeta).getSegmentIndexes).To(
		func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
			return map[int64]*model.SegmentIndex{}
		}).Build()
	defer mockSegmentIndexes.UnPatch()

	snapshotData, err := handler.GenSnapshot(context.Background(), 200)
	require.NoError(t, err)
	require.NotNil(t, snapshotData)
	require.Len(t, snapshotData.Segments, 1)
	assert.Equal(t, int64(2001), snapshotData.Segments[0].GetSegmentId())
	assert.NotEmpty(t, snapshotData.Segments[0].GetManifestPath())
}

func TestGenSnapshot_RejectsSegmentWithoutChannelSeekPosition(t *testing.T) {
	schema := newTestSchema()
	var seg *SegmentInfo

	mockMeta := &meta{indexMeta: &indexMeta{}}
	handler := &ServerHandler{
		s: &Server{
			broker: broker.NewCoordinatorBroker(newMockMixCoord()),
			meta:   mockMeta,
		},
	}

	mockDescribe := mockey.Mock((*mockMixCoord).DescribeCollectionInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
			return &milvuspb.DescribeCollectionResponse{
				Status:              merr.Success(),
				Schema:              schema,
				ShardsNum:           1,
				NumPartitions:       1,
				ConsistencyLevel:    commonpb.ConsistencyLevel_Strong,
				CollectionID:        200,
				VirtualChannelNames: []string{"ch-1"},
			}, nil
		}).Build()
	defer mockDescribe.UnPatch()

	mockShowPartitions := mockey.Mock((*mockMixCoord).ShowPartitionsInternal).To(
		func(m *mockMixCoord, ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
			return &milvuspb.ShowPartitionsResponse{
				Status:         merr.Success(),
				PartitionIDs:   []int64{0},
				PartitionNames: []string{"_default"},
			}, nil
		}).Build()
	defer mockShowPartitions.UnPatch()

	mockSeekPositions := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(
		func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
			return []*msgpb.MsgPosition{
				{ChannelName: "ch-1", Timestamp: 1000, MsgID: []byte{1}},
			}, uint64(1000), nil
		}).Build()
	defer mockSeekPositions.UnPatch()

	mockIndexes := mockey.Mock((*indexMeta).GetIndexesForCollection).To(
		func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
			return []*model.Index{}
		}).Build()
	defer mockIndexes.UnPatch()

	mockSelectSegments := mockey.Mock((*meta).SelectSegments).To(
		func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
			for _, filter := range filters {
				if !filter.Match(seg) {
					return nil
				}
			}
			return []*SegmentInfo{seg}
		}).Build()
	defer mockSelectSegments.UnPatch()

	for _, test := range []struct {
		name          string
		insertChannel string
	}{
		{name: "unknown channel", insertChannel: "ch-missing"},
		{name: "empty channel", insertChannel: ""},
	} {
		t.Run(test.name, func(t *testing.T) {
			seg = NewSegmentInfo(&datapb.SegmentInfo{
				ID:            1001,
				CollectionID:  200,
				PartitionID:   0,
				InsertChannel: test.insertChannel,
				State:         commonpb.SegmentState_Flushed,
				StartPosition: &msgpb.MsgPosition{ChannelName: test.insertChannel, Timestamp: 500},
				DmlPosition:   &msgpb.MsgPosition{ChannelName: test.insertChannel, Timestamp: 600},
				Binlogs: []*datapb.FieldBinlog{
					{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
				},
			})

			snapshotData, err := handler.GenSnapshot(context.Background(), 200)
			require.Error(t, err)
			assert.Nil(t, snapshotData)
			assert.Contains(t, err.Error(), "missing snapshot channel seek position")
		})
	}
}

func TestUncompressIndexFilesPreservesIndexStorePathVersion(t *testing.T) {
	handler := &ServerHandler{
		s: &Server{
			meta: &meta{
				chunkManager: storage.NewLocalChunkManager(objectstorage.RootPath("files")),
				indexMeta:    &indexMeta{},
			},
		},
	}

	mockGetSegmentIndexes := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
		return map[int64]*model.SegmentIndex{
			10: {
				CollectionID:          collectionID,
				PartitionID:           20,
				SegmentID:             segmentID,
				IndexID:               10,
				BuildID:               30,
				IndexVersion:          40,
				IndexState:            commonpb.IndexState_Finished,
				IndexFileKeys:         []string{"index_data"},
				IndexStorePathVersion: indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED,
			},
		}
	}).Build()
	defer mockGetSegmentIndexes.UnPatch()

	mockFieldID := mockey.Mock((*indexMeta).GetFieldIDByIndexID).Return(int64(100)).Build()
	defer mockFieldID.UnPatch()

	mockIndexName := mockey.Mock((*indexMeta).GetIndexNameByID).Return("vec_idx").Build()
	defer mockIndexName.UnPatch()

	mockIndexParams := mockey.Mock((*indexMeta).GetIndexParams).Return(nil).Build()
	defer mockIndexParams.UnPatch()

	mockTypeParams := mockey.Mock((*indexMeta).GetTypeParams).Return(nil).Build()
	defer mockTypeParams.UnPatch()

	indexFiles := uncompressIndexFiles(handler, 1, 2)

	require.Len(t, indexFiles, 1)
	assert.Equal(t, indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED, indexFiles[0].GetIndexStorePathVersion())
	assert.Equal(t, []string{"files/index_v1/1/20/2/30/40/index_data"}, indexFiles[0].GetIndexFilePaths())
}

func TestUncompressJsonStatsForSnapshotPaths(t *testing.T) {
	handler := &ServerHandler{
		s: &Server{
			meta: &meta{
				chunkManager: storage.NewLocalChunkManager(objectstorage.RootPath("files")),
			},
		},
	}

	t.Run("storage v2 path", func(t *testing.T) {
		segInfo := &datapb.SegmentInfo{
			CollectionID: 1,
			PartitionID:  2,
			ID:           3,
		}
		stats := &datapb.JsonKeyStats{
			FieldID:                100,
			BuildID:                10,
			Version:                2,
			JsonKeyStatsDataFormat: 3,
			Files:                  []string{"meta.json", "shared_key_index/.managed.json_0"},
		}

		got := uncompressJSONStats(handler, segInfo, stats)

		assert.Equal(t, []string{
			"files/json_stats/3/10/2/1/2/3/100/meta.json",
			"files/json_stats/3/10/2/1/2/3/100/shared_key_index/.managed.json_0",
		}, got.GetFiles())
		assert.Equal(t, []string{"meta.json", "shared_key_index/.managed.json_0"}, stats.GetFiles())
	})

	t.Run("storage v3 manifest path", func(t *testing.T) {
		segInfo := &datapb.SegmentInfo{
			CollectionID: 1,
			PartitionID:  2,
			ID:           3,
			ManifestPath: `{"ver":7,"base_path":"files/insert_log/1/2/3"}`,
		}
		stats := &datapb.JsonKeyStats{
			FieldID: 100,
			Files:   []string{"meta.json", "shredding_data/data.parquet"},
		}

		got := uncompressJSONStats(handler, segInfo, stats)

		assert.Equal(t, []string{
			"files/insert_log/1/2/3/_stats/json_stats.100/meta.json",
			"files/insert_log/1/2/3/_stats/json_stats.100/shredding_data/data.parquet",
		}, got.GetFiles())
		assert.Equal(t, []string{"meta.json", "shredding_data/data.parquet"}, stats.GetFiles())
	})
}

// TestGenSnapshot_CommitTimestamp verifies that import segments with a
// commit_timestamp are correctly included/excluded based on snapshotTs.
func TestGenSnapshot_CommitTimestamp(t *testing.T) {
	schema := newTestSchema()

	// importSeg has start_position.ts=1000 and commit_timestamp=5000.
	// Its effective timestamp is 5000.
	importSeg := NewSegmentInfo(&datapb.SegmentInfo{
		ID:              999,
		CollectionID:    200,
		PartitionID:     0,
		InsertChannel:   "dml_0_200v0",
		State:           commonpb.SegmentState_Flushed,
		CommitTimestamp: 5000,
		StartPosition:   &msgpb.MsgPosition{Timestamp: 1000},
		Binlogs: []*datapb.FieldBinlog{
			{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
		},
	})

	setupHandlerMocks := func(t *testing.T, snapshotTs uint64) (*ServerHandler, []*mockey.Mocker) {
		mockIndexMeta := &indexMeta{}
		mockMeta := &meta{indexMeta: mockIndexMeta}

		mockMixCoord := mocks2.NewMixCoord(t)
		mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
			Status:              merr.Success(),
			Schema:              schema,
			ShardsNum:           2,
			NumPartitions:       1,
			ConsistencyLevel:    commonpb.ConsistencyLevel_Strong,
			Properties:          []*commonpb.KeyValuePair{},
			CollectionID:        200,
			VirtualChannelNames: []string{"dml_0_200v0"},
		}, nil).Maybe()
		mockMixCoord.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
			Status:         merr.Success(),
			PartitionIDs:   []int64{0},
			PartitionNames: []string{"_default"},
		}, nil).Maybe()

		mockBroker := broker.NewCoordinatorBroker(mockMixCoord)
		handler := &ServerHandler{
			s: &Server{broker: mockBroker, meta: mockMeta},
		}

		var mockers []*mockey.Mocker

		m1 := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
			return []*msgpb.MsgPosition{
				{ChannelName: "dml_0_200v0", Timestamp: snapshotTs, MsgID: []byte{1}},
			}, snapshotTs, nil
		}).Build()
		mockers = append(mockers, m1)

		m2 := mockey.Mock((*indexMeta).GetIndexesForCollection).To(func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
			return []*model.Index{}
		}).Build()
		mockers = append(mockers, m2)

		m3 := mockey.Mock((*meta).SelectSegments).To(func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
			// Apply each filter to the candidate segment list.
			candidates := []*SegmentInfo{importSeg}
			var result []*SegmentInfo
			for _, seg := range candidates {
				pass := true
				for _, f := range filters {
					if !f.Match(seg) {
						pass = false
						break
					}
				}
				if pass {
					result = append(result, seg)
				}
			}
			return result
		}).Build()
		mockers = append(mockers, m3)

		m4 := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
			return nil, false
		}).Build()
		mockers = append(mockers, m4)

		m5 := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
			return map[int64]*model.SegmentIndex{}
		}).Build()
		mockers = append(mockers, m5)

		m6 := mockey.Mock((*meta).GetSegment).To(func(m *meta, ctx context.Context, segmentID int64) *SegmentInfo {
			return nil
		}).Build()
		mockers = append(mockers, m6)

		return handler, mockers
	}

	t.Run("import segment excluded when snapshotTs < commit_timestamp", func(t *testing.T) {
		// snapshotTs=3000 < commit_ts=5000 → segment must NOT appear
		handler, mockers := setupHandlerMocks(t, 3000)
		for _, m := range mockers {
			defer m.UnPatch()
		}

		snapshotData, err := handler.GenSnapshot(context.Background(), 200)
		assert.NoError(t, err)
		assert.NotNil(t, snapshotData)
		assert.Equal(t, 0, len(snapshotData.Segments), "segment with commit_ts=5000 must not appear at snapshotTs=3000")
	})

	t.Run("import segment included when snapshotTs >= commit_timestamp", func(t *testing.T) {
		// snapshotTs=6000 >= commit_ts=5000 → segment must appear
		handler, mockers := setupHandlerMocks(t, 6000)
		for _, m := range mockers {
			defer m.UnPatch()
		}

		snapshotData, err := handler.GenSnapshot(context.Background(), 200)
		assert.NoError(t, err)
		assert.NotNil(t, snapshotData)
		assert.Equal(t, 1, len(snapshotData.Segments), "segment with commit_ts=5000 must appear at snapshotTs=6000")
		assert.Equal(t, int64(999), snapshotData.Segments[0].SegmentId)
	})
}
