// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package sidx

import (
	"errors"
	"path/filepath"
	"sort"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/stretchr/testify/require"

	"github.com/apache/skywalking-banyandb/api/common"
	"github.com/apache/skywalking-banyandb/banyand/protector"
	"github.com/apache/skywalking-banyandb/pkg/convert"
	"github.com/apache/skywalking-banyandb/pkg/fs"
	pbv1 "github.com/apache/skywalking-banyandb/pkg/pb/v1"
	"github.com/apache/skywalking-banyandb/pkg/test"
)

var conventionalBlock = block{
	userKeys: []int64{1, 2},
	data:     [][]byte{[]byte("data1"), []byte("data2")},
	tags: map[string]*tagData{
		"service": {
			name:      "service",
			valueType: pbv1.ValueTypeStr,
			values: []tagRow{
				{value: []byte("service1")},
				{value: []byte("service2")},
			},
		},
	},
}

var mergedBlock = block{
	userKeys: []int64{1, 2, 3, 4},
	data:     [][]byte{[]byte("data1"), []byte("data2"), []byte("data3"), []byte("data4")},
	tags: map[string]*tagData{
		"arrTag": {
			name:      "arrTag",
			valueType: pbv1.ValueTypeStrArr,
			values: []tagRow{
				{valueArr: [][]byte{[]byte("value1"), []byte("value2")}},
				{valueArr: [][]byte{[]byte("value3"), []byte("value4")}},
				{valueArr: [][]byte{[]byte("value5"), []byte("value6")}},
				{valueArr: [][]byte{[]byte("value7"), []byte("value8")}},
			},
		},
	},
}

var duplicatedMergedBlock = block{
	userKeys: []int64{1, 2, 2, 3, 3, 4},
	data:     [][]byte{[]byte("data1"), []byte("data2"), []byte("data3"), []byte("data5"), []byte("data4"), []byte("data6")},
	tags: map[string]*tagData{
		"arrTag": {
			name:      "arrTag",
			valueType: pbv1.ValueTypeStrArr,
			values: []tagRow{
				{valueArr: [][]byte{[]byte("value1"), []byte("value2")}},
				{valueArr: [][]byte{[]byte("duplicated1")}},
				{valueArr: [][]byte{[]byte("value3"), []byte("value4")}},
				{valueArr: [][]byte{[]byte("value5"), []byte("value6")}},
				{valueArr: [][]byte{[]byte("duplicated2")}},
				{valueArr: [][]byte{[]byte("value7"), []byte("value8")}},
			},
		},
	},
}

func Test_mergeTwoBlocks(t *testing.T) {
	tests := []struct {
		left  *blockPointer
		right *blockPointer
		want  *blockPointer
		name  string
	}{
		{
			name:  "Merge two empty blocks",
			left:  &blockPointer{},
			right: &blockPointer{},
			want:  &blockPointer{},
		},
		{
			name:  "Merge left is non-empty right is empty",
			left:  &blockPointer{block: deepCopyBlock(&conventionalBlock)},
			right: &blockPointer{},
			want:  &blockPointer{block: deepCopyBlock(&conventionalBlock), bm: blockMetadata{minKey: 1, maxKey: 2}},
		},
		{
			name:  "Merge left is empty right is non-empty",
			left:  &blockPointer{},
			right: &blockPointer{block: deepCopyBlock(&conventionalBlock)},
			want:  &blockPointer{block: deepCopyBlock(&conventionalBlock), bm: blockMetadata{minKey: 1, maxKey: 2}},
		},
		{
			name: "Merge two non-empty blocks without overlap",
			left: &blockPointer{
				block: block{
					userKeys: []int64{1, 2},
					data:     [][]byte{[]byte("data1"), []byte("data2")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value1"), []byte("value2")}},
								{valueArr: [][]byte{[]byte("value3"), []byte("value4")}},
							},
						},
					},
				},
			},
			right: &blockPointer{
				block: block{
					userKeys: []int64{3, 4},
					data:     [][]byte{[]byte("data3"), []byte("data4")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value5"), []byte("value6")}},
								{valueArr: [][]byte{[]byte("value7"), []byte("value8")}},
							},
						},
					},
				},
			},
			want: &blockPointer{block: deepCopyBlock(&mergedBlock), bm: blockMetadata{minKey: 1, maxKey: 4}},
		},
		{
			name: "Merge two non-empty blocks without duplicated userKeys",
			left: &blockPointer{
				block: block{
					userKeys: []int64{1, 3},
					data:     [][]byte{[]byte("data1"), []byte("data3")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value1"), []byte("value2")}},
								{valueArr: [][]byte{[]byte("value5"), []byte("value6")}},
							},
						},
					},
				},
			},
			right: &blockPointer{
				block: block{
					userKeys: []int64{2, 4},
					data:     [][]byte{[]byte("data2"), []byte("data4")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value3"), []byte("value4")}},
								{valueArr: [][]byte{[]byte("value7"), []byte("value8")}},
							},
						},
					},
				},
			},
			want: &blockPointer{block: deepCopyBlock(&mergedBlock), bm: blockMetadata{minKey: 1, maxKey: 4}},
		},
		{
			name: "Merge two non-empty blocks with duplicated userKeys",
			left: &blockPointer{
				block: block{
					userKeys: []int64{1, 2, 3},
					data:     [][]byte{[]byte("data1"), []byte("data2"), []byte("data4")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value1"), []byte("value2")}},
								{valueArr: [][]byte{[]byte("duplicated1")}},
								{valueArr: [][]byte{[]byte("duplicated2")}},
							},
						},
					},
				},
			},
			right: &blockPointer{
				block: block{
					userKeys: []int64{2, 3, 4},
					data:     [][]byte{[]byte("data3"), []byte("data5"), []byte("data6")},
					tags: map[string]*tagData{
						"arrTag": {
							name:      "arrTag",
							valueType: pbv1.ValueTypeStrArr,
							values: []tagRow{
								{valueArr: [][]byte{[]byte("value3"), []byte("value4")}},
								{valueArr: [][]byte{[]byte("value5"), []byte("value6")}},
								{valueArr: [][]byte{[]byte("value7"), []byte("value8")}},
							},
						},
					},
				},
			},
			want: &blockPointer{block: deepCopyBlock(&duplicatedMergedBlock), bm: blockMetadata{minKey: 1, maxKey: 4}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			target := &blockPointer{}
			mergeTwoBlocks(target, tt.left, tt.right)
			if diff := cmp.Diff(target, tt.want,
				cmpopts.IgnoreFields(tagData{}, "uniqueValues", "tmpBytes"),
				cmpopts.IgnoreFields(blockMetadata{}, "tagsBlocks", "tagProjection", "dataBlock", "keysBlock", "seriesID", "uncompressedSize", "count", "keysEncodeType"),
				cmp.AllowUnexported(block{}, tagData{}, tagRow{}, blockPointer{}, blockMetadata{}),
			); diff != "" {
				t.Errorf("mergeTwoBlocks() mismatch (-got +want):\n%s", diff)
			}
		})
	}
}

func generateHugeElements(start, end int64, seriesIDBase common.SeriesID) *elements {
	es := generateElements()
	for i := start; i <= end; i++ {
		seriesID := seriesIDBase
		if i%1000 == 0 {
			seriesID = seriesIDBase + 1
		}
		if i%2000 == 0 {
			seriesID = seriesIDBase + 2
		}

		tags := []Tag{
			{Name: "service", Value: []byte("test-service"), ValueType: pbv1.ValueTypeStr},
		}
		data := make([]byte, 50)
		es.mustAppend(seriesID, i, data, tags)
	}
	return es
}

var (
	es1 = func() *elements {
		es := generateElements()
		es.mustAppend(1, 100, make([]byte, 1600), []Tag{
			{Name: "service", Value: []byte("service1"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(2, 200, make([]byte, 40), []Tag{
			{Name: "env", Value: []byte("prod"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(3, 300, make([]byte, 25), []Tag{
			{Name: "region", Value: []byte("us"), ValueType: pbv1.ValueTypeStr},
		})
		return es
	}()

	es2 = func() *elements {
		es := generateElements()
		es.mustAppend(1, 150, make([]byte, 1600), []Tag{
			{Name: "service", Value: []byte("service1"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(2, 250, make([]byte, 40), []Tag{
			{Name: "env", Value: []byte("prod"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(3, 350, make([]byte, 25), []Tag{
			{Name: "region", Value: []byte("us"), ValueType: pbv1.ValueTypeStr},
		})
		return es
	}()

	esStrArr1 = func() *elements {
		es := generateElements()
		es.mustAppend(1, 100, make([]byte, 100), []Tag{
			{Name: "arrTag", ValueArr: [][]byte{[]byte("value1"), []byte("value2")}, ValueType: pbv1.ValueTypeStrArr},
		})
		es.mustAppend(2, 200, make([]byte, 100), []Tag{
			{Name: "arrTag", ValueArr: [][]byte{[]byte("value3"), []byte("value4")}, ValueType: pbv1.ValueTypeStrArr},
		})
		return es
	}()

	esStrArr2 = func() *elements {
		es := generateElements()
		es.mustAppend(1, 150, make([]byte, 100), []Tag{
			{Name: "arrTag", ValueArr: [][]byte{[]byte("value5"), []byte("value6")}, ValueType: pbv1.ValueTypeStrArr},
		})
		es.mustAppend(2, 250, make([]byte, 100), []Tag{
			{Name: "arrTag", ValueArr: [][]byte{[]byte("value7"), []byte("value8")}, ValueType: pbv1.ValueTypeStrArr},
		})
		return es
	}()

	esStrArrWithEmpty = func() *elements {
		es := generateElements()
		es.mustAppend(1, 300, make([]byte, 100), []Tag{
			{Name: "arrTag", ValueArr: [][]byte{[]byte("a"), []byte(""), []byte("b")}, ValueType: pbv1.ValueTypeStrArr},
		})
		return es
	}()
)

func Test_mergeParts(t *testing.T) {
	tests := []struct {
		wantErr error
		name    string
		esList  []*elements
		want    []blockMetadata
	}{
		{
			name:    "Test with no element",
			esList:  []*elements{},
			wantErr: errNoPartToMerge,
		},
		{
			name:   "Test with single part",
			esList: []*elements{es1},
			want: []blockMetadata{
				{seriesID: 1, count: 1, uncompressedSize: 1623},
				{seriesID: 2, count: 1, uncompressedSize: 55},
				{seriesID: 3, count: 1, uncompressedSize: 41},
			},
		},
		{
			name:   "Test with multiple parts with different userKeys",
			esList: []*elements{es1, es2, es2},
			want: []blockMetadata{
				{seriesID: 1, count: 3, uncompressedSize: 4869},
				{seriesID: 2, count: 3, uncompressedSize: 165},
				{seriesID: 3, count: 3, uncompressedSize: 123},
			},
		},
		{
			name:   "Test with multiple parts with same userKeys",
			esList: []*elements{es1, es1, es1},
			want: []blockMetadata{
				{seriesID: 1, count: 3, uncompressedSize: 4869},
				{seriesID: 2, count: 3, uncompressedSize: 165},
				{seriesID: 3, count: 3, uncompressedSize: 123},
			},
		},
		{
			name:   "Test with multiple parts with a large quantity of different userKeys",
			esList: []*elements{generateHugeElements(1, 1000, 1), generateHugeElements(1001, 2000, 2)},
			want: []blockMetadata{
				{seriesID: 1, count: 999, uncompressedSize: 76923},
				{seriesID: 2, count: 1000, uncompressedSize: 77000},
				{seriesID: 4, count: 1, uncompressedSize: 77},
			},
		},
		{
			name:   "Test with string array value type",
			esList: []*elements{esStrArr1, esStrArr2},
			want: []blockMetadata{
				{seriesID: 1, count: 2, uncompressedSize: 264},
				{seriesID: 2, count: 2, uncompressedSize: 264},
			},
		},
		{
			name:   "Test with string array containing empty strings",
			esList: []*elements{esStrArrWithEmpty},
			want: []blockMetadata{
				{seriesID: 1, count: 1, uncompressedSize: 128},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			wantBlocks := elementsToBlocks(tt.esList)

			verify := func(t *testing.T, pp []*partWrapper, fileSystem fs.FileSystem, root string, partID uint64) {
				closeCh := make(chan struct{})
				defer close(closeCh)
				s := &sidx{pm: protector.Nop{}}
				p, err := s.mergeParts(fileSystem, closeCh, pp, partID, root, nil)
				if tt.wantErr != nil {
					if !errors.Is(err, tt.wantErr) {
						t.Fatalf("Unexpected error: got %v, want %v", err, tt.wantErr)
					}
					return
				}
				defer func() {
					if p != nil {
						p.release()
					}
				}()
				pmi := &partMergeIter{}
				pmi.mustInitFromPart(p.p)
				reader := &blockReader{}
				reader.init([]*partMergeIter{pmi})
				var got []blockMetadata
				var gotBlocks []block
				decoder := generateTagValuesDecoder()
				defer releaseTagValuesDecoder(decoder)

				for reader.nextBlockMetadata() {
					got = append(got, reader.block.bm)
					reader.loadBlockData(decoder)
					gotBlocks = append(gotBlocks, deepCopyBlock(&reader.block.block))
				}
				require.NoError(t, reader.error())

				if diff := cmp.Diff(got, tt.want,
					cmpopts.IgnoreFields(blockMetadata{}, "tagsBlocks"),
					cmpopts.IgnoreFields(blockMetadata{}, "tagProjection"),
					cmpopts.IgnoreFields(blockMetadata{}, "dataBlock"),
					cmpopts.IgnoreFields(blockMetadata{}, "keysBlock"),
					cmpopts.IgnoreFields(blockMetadata{}, "minKey"),
					cmpopts.IgnoreFields(blockMetadata{}, "maxKey"),
					cmpopts.IgnoreFields(blockMetadata{}, "keysEncodeType"),
					cmp.AllowUnexported(blockMetadata{}),
				); diff != "" {
					t.Errorf("Unexpected blockMetadata (-got +want):\n%s", diff)
				}

				if diff := cmp.Diff(gotBlocks, wantBlocks,
					cmpopts.IgnoreFields(tagData{}, "uniqueValues", "tmpBytes"),
					cmp.AllowUnexported(block{}, tagData{}, tagRow{}),
				); diff != "" {
					t.Errorf("Unexpected blocks (-got +want):\n%s", diff)
				}
			}

			t.Run("memory parts", func(t *testing.T) {
				var pp []*partWrapper
				tmpPath, defFn := test.Space(require.New(t))
				defer func() {
					for _, pw := range pp {
						pw.release()
					}
					defFn()
				}()
				for _, es := range tt.esList {
					mp := GenerateMemPart()
					mp.mustInitFromElements(es)
					pp = append(pp, newPartWrapper(mp, openMemPart(mp)))
				}
				verify(t, pp, fs.NewLocalFileSystem(), tmpPath, 1)
			})

			t.Run("file parts", func(t *testing.T) {
				var fpp []*partWrapper
				tmpPath, defFn := test.Space(require.New(t))
				defer func() {
					for _, pw := range fpp {
						pw.release()
					}
					defFn()
				}()
				fileSystem := fs.NewLocalFileSystem()
				for i, es := range tt.esList {
					mp := GenerateMemPart()
					mp.mustInitFromElements(es)
					partPath := filepath.Join(tmpPath, "part_"+string(rune('0'+i)))
					mp.mustFlush(fileSystem, partPath)
					filePart := mustOpenPart(uint64(i), partPath, fileSystem)
					filePW := newPartWrapper(nil, filePart)
					fpp = append(fpp, filePW)
					ReleaseMemPart(mp)
				}
				verify(t, fpp, fileSystem, tmpPath, uint64(len(tt.esList)))
			})
		})
	}
}

func Test_mergePartsWithConflictTags(t *testing.T) {
	esConflictStr := func() *elements {
		es := generateElements()
		es.mustAppend(1, 100, make([]byte, 10), []Tag{
			{Name: "status", Value: []byte("ok"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(2, 200, make([]byte, 10), []Tag{
			{Name: "status", Value: []byte("error"), ValueType: pbv1.ValueTypeStr},
		})
		es.mustAppend(3, 300, make([]byte, 10), []Tag{
			{Name: "status", Value: []byte("warn"), ValueType: pbv1.ValueTypeStr},
		})
		return es
	}()

	esConflictInt := func() *elements {
		es := generateElements()
		es.mustAppend(1, 150, make([]byte, 10), []Tag{
			{Name: "status", Value: convert.Int64ToBytes(200), ValueType: pbv1.ValueTypeInt64},
		})
		es.mustAppend(2, 250, make([]byte, 10), []Tag{
			{Name: "status", Value: convert.Int64ToBytes(500), ValueType: pbv1.ValueTypeInt64},
		})
		es.mustAppend(3, 350, make([]byte, 10), []Tag{
			{Name: "status", Value: convert.Int64ToBytes(404), ValueType: pbv1.ValueTypeInt64},
		})
		return es
	}()

	verify := func(t *testing.T, pp []*partWrapper, fileSystem fs.FileSystem, root string, partID uint64) {
		closeCh := make(chan struct{})
		defer close(closeCh)
		s := &sidx{pm: protector.Nop{}}
		p, mergeErr := s.mergeParts(fileSystem, closeCh, pp, partID, root, nil)
		require.NoError(t, mergeErr)
		defer func() {
			if p != nil {
				p.release()
			}
		}()

		pmi := &partMergeIter{}
		pmi.mustInitFromPart(p.p)
		reader := &blockReader{}
		reader.init([]*partMergeIter{pmi})
		decoder := generateTagValuesDecoder()
		defer releaseTagValuesDecoder(decoder)

		var blockCount int
		for reader.nextBlockMetadata() {
			reader.loadBlockData(decoder)
			b := &reader.block.block
			blockCount++

			_, hasBareStatus := b.tags["status"]
			require.False(t, hasBareStatus,
				"block %d (seriesID=%d) has bare 'status' tag; expected renamed tags only",
				blockCount, reader.block.bm.seriesID)

			hasStr := false
			hasInt := false
			for tagName := range b.tags {
				if tagName == "status#str" {
					hasStr = true
				}
				if tagName == "status#int" {
					hasInt = true
				}
			}
			require.True(t, hasStr || hasInt,
				"block %d (seriesID=%d) has neither 'status#str' nor 'status#int'",
				blockCount, reader.block.bm.seriesID)
		}
		require.NoError(t, reader.error())
		require.Equal(t, 3, blockCount, "expected 3 blocks (one per seriesID)")
	}

	esList := []*elements{esConflictStr, esConflictInt}

	t.Run("memory parts", func(t *testing.T) {
		var pp []*partWrapper
		tmpPath, defFn := test.Space(require.New(t))
		defer func() {
			for _, pw := range pp {
				pw.release()
			}
			defFn()
		}()
		for _, es := range esList {
			mp := GenerateMemPart()
			mp.mustInitFromElements(es)
			pp = append(pp, newPartWrapper(mp, openMemPart(mp)))
		}
		verify(t, pp, fs.NewLocalFileSystem(), tmpPath, 1)
	})

	t.Run("file parts", func(t *testing.T) {
		var fpp []*partWrapper
		tmpPath, defFn := test.Space(require.New(t))
		defer func() {
			for _, pw := range fpp {
				pw.release()
			}
			defFn()
		}()
		fileSystem := fs.NewLocalFileSystem()
		for i, es := range esList {
			mp := GenerateMemPart()
			mp.mustInitFromElements(es)
			partPath := filepath.Join(tmpPath, "part_"+string(rune('0'+i)))
			mp.mustFlush(fileSystem, partPath)
			filePart := mustOpenPart(uint64(i), partPath, fileSystem)
			filePW := newPartWrapper(nil, filePart)
			fpp = append(fpp, filePW)
			ReleaseMemPart(mp)
		}
		verify(t, fpp, fileSystem, tmpPath, uint64(len(esList)))
	})
}

func elementsToBlocks(esList []*elements) []block {
	merged := generateElements()
	defer releaseElements(merged)

	for _, es := range esList {
		for i := 0; i < len(es.seriesIDs); i++ {
			var tags []Tag
			for _, t := range es.tags[i] {
				tags = append(tags, Tag{
					Name:      t.name,
					Value:     t.value,
					ValueArr:  t.valueArr,
					ValueType: t.valueType,
				})
			}
			merged.mustAppend(es.seriesIDs[i], es.userKeys[i], es.data[i], tags)
		}
	}

	sort.Sort(merged)

	var blocks []block
	if merged.Len() == 0 {
		return blocks
	}

	start := 0
	for i := 1; i <= merged.Len(); i++ {
		if i == merged.Len() || merged.seriesIDs[i] != merged.seriesIDs[start] {
			b := block{
				tags:     make(map[string]*tagData),
				userKeys: make([]int64, i-start),
				data:     make([][]byte, i-start),
			}
			copy(b.userKeys, merged.userKeys[start:i])
			for k := 0; k < i-start; k++ {
				b.data[k] = merged.data[start+k]
			}
			(&b).mustInitFromTags(merged.tags[start:i])
			blocks = append(blocks, b)
			start = i
		}
	}
	return blocks
}

func deepCopyBlock(b *block) block {
	newB := block{
		tags: make(map[string]*tagData),
	}
	newB.userKeys = append([]int64(nil), b.userKeys...)
	newB.data = make([][]byte, len(b.data))
	for i, d := range b.data {
		newB.data[i] = cloneBytes(d)
	}
	for k, v := range b.tags {
		newTd := &tagData{
			name:      v.name,
			valueType: v.valueType,
			values:    make([]tagRow, len(v.values)),
		}
		for i, row := range v.values {
			newRow := tagRow{}
			if row.value != nil {
				newRow.value = cloneBytes(row.value)
			}
			if row.valueArr != nil {
				newRow.valueArr = make([][]byte, len(row.valueArr))
				for j, arrVal := range row.valueArr {
					newRow.valueArr[j] = cloneBytes(arrVal)
				}
			}
			newTd.values[i] = newRow
		}
		newB.tags[k] = newTd
	}
	return newB
}

func cloneBytes(src []byte) []byte {
	if src == nil {
		return nil
	}
	dst := make([]byte, len(src))
	copy(dst, src)
	return dst
}

func Test_filterBlockPointer_withKeepPredicate(t *testing.T) {
	// 5 entries: rows 0,2,4 start with 'k' (kept); rows 1,3 start with 'd' (dropped).
	bp := &blockPointer{
		block: block{
			userKeys: []int64{10, 20, 30, 40, 50},
			data: [][]byte{
				{'k', '1'}, {'d', '2'}, {'k', '3'}, {'d', '4'}, {'k', '5'},
			},
			tags: map[string]*tagData{
				"svc": {
					name:      "svc",
					valueType: pbv1.ValueTypeStr,
					values: []tagRow{
						{value: []byte("svc1")},
						{value: []byte("svc2")},
						{value: []byte("svc3")},
						{value: []byte("svc4")},
						{value: []byte("svc5")},
					},
				},
			},
		},
	}

	filterBlockPointer(bp, func(d []byte) bool { return len(d) > 0 && d[0] == 'k' })

	require.Equal(t, []int64{10, 30, 50}, bp.userKeys, "userKeys row-alignment")
	require.Equal(t, [][]byte{{'k', '1'}, {'k', '3'}, {'k', '5'}}, bp.data, "data row-alignment")
	svcVals := bp.tags["svc"].values
	require.Len(t, svcVals, 3, "tag values count after filter")
	require.Equal(t, []byte("svc1"), svcVals[0].value)
	require.Equal(t, []byte("svc3"), svcVals[1].value)
	require.Equal(t, []byte("svc5"), svcVals[2].value)
	require.Equal(t, int64(10), bp.bm.minKey, "minKey after updateMetadata")
	require.Equal(t, int64(50), bp.bm.maxKey, "maxKey after updateMetadata")
}

func Test_mergeParts_withKeepPredicate(t *testing.T) {
	// Encode trace IDs using the same byte layout as the trace engine: [0x01] + traceID.
	encodeID := func(id string) []byte { return append([]byte{0x01}, []byte(id)...) }

	es := generateElements()
	es.mustAppend(1, 100, encodeID("traceA"), []Tag{{Name: "svc", Value: []byte("a"), ValueType: pbv1.ValueTypeStr}})
	es.mustAppend(1, 200, encodeID("traceB"), []Tag{{Name: "svc", Value: []byte("b"), ValueType: pbv1.ValueTypeStr}})
	es.mustAppend(1, 300, encodeID("traceC"), []Tag{{Name: "svc", Value: []byte("c"), ValueType: pbv1.ValueTypeStr}})

	keep := func(data []byte) bool {
		if len(data) < 2 || data[0] != 0x01 {
			return true
		}
		return string(data[1:]) != "traceB"
	}

	tmpPath, defFn := test.Space(require.New(t))
	defer defFn()
	fileSystem := fs.NewLocalFileSystem()

	mp := GenerateMemPart()
	mp.mustInitFromElements(es)
	pw := newPartWrapper(mp, openMemPart(mp))
	defer pw.release()

	closeCh := make(chan struct{})
	defer close(closeCh)
	s := &sidx{pm: protector.Nop{}}
	merged, mergeErr := s.mergeParts(fileSystem, closeCh, []*partWrapper{pw}, 99, tmpPath, keep)
	require.NoError(t, mergeErr)
	defer merged.release()

	pmi := &partMergeIter{}
	pmi.mustInitFromPart(merged.p)
	reader := &blockReader{}
	reader.init([]*partMergeIter{pmi})
	decoder := generateTagValuesDecoder()
	defer releaseTagValuesDecoder(decoder)

	var retained []string
	for reader.nextBlockMetadata() {
		reader.loadBlockData(decoder)
		for _, d := range reader.block.block.data {
			if len(d) >= 2 && d[0] == 0x01 {
				retained = append(retained, string(d[1:]))
			}
		}
	}
	require.NoError(t, reader.error())

	require.NotContains(t, retained, "traceB", "traceB must be absent from merged part")
	require.Len(t, retained, 2, "exactly traceA and traceC retained")
	require.Equal(t, uint64(2), merged.p.partMetadata.TotalCount, "partMetadata TotalCount == retained count")
}
