// 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 (
	"context"
	"fmt"
	"sort"
	"strings"
	"testing"
	"time"

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

	"github.com/apache/skywalking-banyandb/api/common"
	modelv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/model/v1"
	"github.com/apache/skywalking-banyandb/banyand/internal/test"
	"github.com/apache/skywalking-banyandb/pkg/index"
	pbv1 "github.com/apache/skywalking-banyandb/pkg/pb/v1"
	"github.com/apache/skywalking-banyandb/pkg/query/model"
)

type strictQuotaProtector struct {
	test.MockMemoryProtector
}

func (s *strictQuotaProtector) AvailableBytes() int64 {
	return 0
}

func TestSIDX_Query_BasicQuery(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 100, "data1"),
		createTestWriteRequest(1, 101, "data2"),
	}
	writeTestData(t, sidx, reqs, 4, 4)

	waitForIntroducerLoop()

	queryReq := createTestQueryRequest(1)
	resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

	var keys []int64
	for res := range resultsCh {
		require.NoError(t, res.Error)
		keys = append(keys, res.Keys...)
	}
	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	if len(keys) > 0 {
		assert.Greater(t, len(keys), 0)
	}
}

func TestSIDX_QuerySyncMatchesStreamingQuery(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()
	reqs := []WriteRequest{
		createTestWriteRequest(1, 100, "data100"),
		createTestWriteRequest(1, 200, "data200"),
		createTestWriteRequest(1, 150, "data150"),
	}
	writeTestData(t, sidx, reqs, 44, 44)
	waitForIntroducerLoop()

	queryReq := createTestQueryRequest(1)
	queryReq.MaxBatchSize = 2
	streamingRows := collectStreamingRows(ctx, t, sidx, queryReq)
	syncQuerier, ok := sidx.(interface {
		QuerySync(context.Context, QueryRequest) ([]*QueryResponse, error)
	})
	require.True(t, ok)
	syncRows, err := syncQuerier.QuerySync(ctx, queryReq)
	require.NoError(t, err)
	require.Equal(t, flattenQueryRows(streamingRows), flattenQueryRows(syncRows))
}

// TestSIDX_QuerySync_EarlyStopAtMaxBatchSize verifies the sync query path stops
// scanning once MaxBatchSize distinct results are collected, while still returning
// the correct ordered top-N. Blocks are per series, so writing many single-element
// series forces the scan across multiple batches (blockScannerBatchSize); the block
// iterator yields them in ascending key order, so the bounded result's prefix must
// equal the uncapped, fully-ordered scan's prefix.
func TestSIDX_QuerySync_EarlyStopAtMaxBatchSize(t *testing.T) {
	s := createTestSIDX(t)
	defer func() {
		assert.NoError(t, s.Close())
	}()
	ctx := context.Background()

	const seriesCount = 80
	seriesIDs := make([]common.SeriesID, 0, seriesCount)
	for i := range seriesCount {
		key := int64(i + 1)
		series := common.SeriesID(i + 1)
		seriesIDs = append(seriesIDs, series)
		writeTestData(t, s, []WriteRequest{
			createTestWriteRequest(series, key, fmt.Sprintf("data%03d", key)),
		}, 44, uint64(i+1))
	}
	waitForIntroducerLoop()

	syncQuerier, ok := s.(interface {
		QuerySync(context.Context, QueryRequest) ([]*QueryResponse, error)
	})
	require.True(t, ok)

	baseReq := createTestQueryRequest(seriesIDs...)

	// Uncapped scan reads every block and returns all elements in key order.
	uncapped, err := syncQuerier.QuerySync(ctx, baseReq)
	require.NoError(t, err)
	uncappedRows := flattenQueryRows(uncapped)
	require.Len(t, uncappedRows, seriesCount)

	const maxBatch = 10
	cappedReq := baseReq
	cappedReq.MaxBatchSize = maxBatch
	capped, err := syncQuerier.QuerySync(ctx, cappedReq)
	require.NoError(t, err)
	cappedRows := flattenQueryRows(capped)

	// Bounded query collected at least MaxBatchSize and stopped before reading every block.
	require.GreaterOrEqual(t, len(cappedRows), maxBatch, "must collect at least MaxBatchSize results")
	require.Less(t, len(cappedRows), seriesCount, "must stop scanning before reading every block")

	// The bounded result is the correct ordered top-N: its first MaxBatchSize rows
	// equal the first MaxBatchSize rows of the uncapped, fully-ordered scan.
	require.Equal(t, uncappedRows[:maxBatch], cappedRows[:maxBatch])
}

func TestSIDX_Query_EmptyResult(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	queryReq := createTestQueryRequest(999)
	resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

	resultCount := 0
	for res := range resultsCh {
		require.NoError(t, res.Error)
		resultCount += res.Len()
	}
	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	assert.Equal(t, 0, resultCount)
}

type queryRow struct {
	data   string
	key    int64
	series common.SeriesID
	partID uint64
}

func collectStreamingRows(ctx context.Context, t *testing.T, sidx SIDX, req QueryRequest) []*QueryResponse {
	t.Helper()
	resultsCh, errCh := sidx.StreamingQuery(ctx, req)
	var results []*QueryResponse
	for res := range resultsCh {
		require.NoError(t, res.Error)
		results = append(results, res)
	}
	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}
	return results
}

func flattenQueryRows(results []*QueryResponse) []queryRow {
	var rows []queryRow
	for _, result := range results {
		for rowIdx := range result.Keys {
			rows = append(rows, queryRow{
				key:    result.Keys[rowIdx],
				data:   string(result.Data[rowIdx]),
				series: result.SIDs[rowIdx],
				partID: result.PartIDs[rowIdx],
			})
		}
	}
	return rows
}

func TestSIDX_Query_KeyRangeFilter(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 100, "data100"),
		createTestWriteRequest(1, 150, "data150"),
		createTestWriteRequest(1, 200, "data200"),
	}
	writeTestData(t, sidx, reqs, 5, 5)

	waitForIntroducerLoop()

	minKey := int64(120)
	maxKey := int64(180)
	queryReq := QueryRequest{
		SeriesIDs: []common.SeriesID{1},
		MinKey:    &minKey,
		MaxKey:    &maxKey,
	}

	resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

	var keys []int64
	for res := range resultsCh {
		require.NoError(t, res.Error)
		keys = append(keys, res.Keys...)
	}
	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	if len(keys) > 0 {
		for _, key := range keys {
			assert.GreaterOrEqual(t, key, minKey)
			assert.LessOrEqual(t, key, maxKey)
		}
	}
}

func TestSIDX_Query_Ordering(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 300, "series1-data300"),
		createTestWriteRequest(1, 100, "series1-data100"),
		createTestWriteRequest(1, 200, "series1-data200"),
		createTestWriteRequest(2, 250, "series2-data250"),
		createTestWriteRequest(2, 150, "series2-data150"),
		createTestWriteRequest(2, 50, "series2-data50"),
		createTestWriteRequest(3, 350, "series3-data350"),
		createTestWriteRequest(3, 75, "series3-data75"),
		createTestWriteRequest(3, 175, "series3-data175"),
	}
	writeTestData(t, sidx, reqs, 6, 6)

	waitForIntroducerLoop()

	tests := []struct {
		order     *index.OrderBy
		name      string
		seriesIDs []common.SeriesID
		ascending bool
	}{
		{
			name:      "ascending order single series",
			ascending: true,
			order:     &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			seriesIDs: []common.SeriesID{1},
		},
		{
			name:      "descending order single series",
			ascending: false,
			order:     &index.OrderBy{Sort: modelv1.Sort_SORT_DESC},
			seriesIDs: []common.SeriesID{1},
		},
		{
			name:      "ascending order multiple series",
			ascending: true,
			order:     &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			seriesIDs: []common.SeriesID{1, 2, 3},
		},
		{
			name:      "descending order multiple series",
			ascending: false,
			order:     &index.OrderBy{Sort: modelv1.Sort_SORT_DESC},
			seriesIDs: []common.SeriesID{1, 2, 3},
		},
		{
			name:      "default order multiple series",
			ascending: true,
			order:     nil,
			seriesIDs: []common.SeriesID{1, 2, 3},
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			queryReq := QueryRequest{
				SeriesIDs: tt.seriesIDs,
				Order:     tt.order,
			}

			resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

			var allKeys []int64
			var allData [][]byte
			var allSIDs []common.SeriesID
			for res := range resultsCh {
				require.NoError(t, res.Error)
				allKeys = append(allKeys, res.Keys...)
				allData = append(allData, res.Data...)
				allSIDs = append(allSIDs, res.SIDs...)
			}
			if err, ok := <-errCh; ok {
				require.NoError(t, err)
			}

			t.Logf("Query with series %v, order %v", tt.seriesIDs, tt.order)
			for i, key := range allKeys {
				if i < len(allData) && i < len(allSIDs) {
					t.Logf("  Key: %d, Data: %s, SeriesID: %d", key, string(allData[i]), allSIDs[i])
				}
			}

			if len(allKeys) > 1 {
				if len(tt.seriesIDs) == 1 {
					// For single series, verify global sorting
					isSorted := sort.SliceIsSorted(allKeys, func(i, j int) bool {
						if tt.ascending {
							return allKeys[i] < allKeys[j]
						}
						return allKeys[i] > allKeys[j]
					})
					assert.True(t, isSorted, "Keys should be sorted in %s order. Keys: %v",
						map[bool]string{true: "ascending", false: "descending"}[tt.ascending], allKeys)
				} else {
					// For multiple series, verify sorting within each series group
					seriesGroups := make(map[common.SeriesID][]int64)
					for i, sid := range allSIDs {
						if i < len(allKeys) {
							seriesGroups[sid] = append(seriesGroups[sid], allKeys[i])
						}
					}
					for sid, keys := range seriesGroups {
						if len(keys) > 1 {
							isSorted := sort.SliceIsSorted(keys, func(i, j int) bool {
								if tt.ascending {
									return keys[i] < keys[j]
								}
								return keys[i] > keys[j]
							})
							assert.True(t, isSorted, "Keys for series %d should be sorted in %s order. Keys: %v",
								sid, map[bool]string{true: "ascending", false: "descending"}[tt.ascending], keys)
						}
					}
				}
			}
		})
	}
}

func TestSIDX_Query_WithArrValues(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 100, "data100", Tag{
			Name: "arr_tag",
			ValueArr: [][]byte{
				[]byte("a"),
				[]byte("b"),
			},
			ValueType: pbv1.ValueTypeStrArr,
		}),
		createTestWriteRequest(1, 150, "data150"),
		createTestWriteRequest(1, 200, "data200"),
	}
	writeTestData(t, sidx, reqs, 7, 7)

	waitForIntroducerLoop()

	queryReq := QueryRequest{
		SeriesIDs: []common.SeriesID{1},
		TagProjection: []model.TagProjection{
			{
				Names: []string{"arr_tag"},
			},
		},
	}

	resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

	var keys []int64
	for res := range resultsCh {
		require.NoError(t, res.Error)
		keys = append(keys, res.Keys...)
	}
	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	assert.Equal(t, 3, len(keys))
}

func TestSIDX_Query_Validation(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	tests := []struct {
		name      string
		req       QueryRequest
		expectErr bool
	}{
		{
			name:      "valid request",
			req:       createTestQueryRequest(1),
			expectErr: false,
		},
		{
			name:      "empty series IDs",
			req:       QueryRequest{SeriesIDs: []common.SeriesID{}},
			expectErr: true,
		},
		{
			name: "invalid key range",
			req: QueryRequest{
				SeriesIDs: []common.SeriesID{1},
				MinKey:    ptrInt64(200),
				MaxKey:    ptrInt64(100),
			},
			expectErr: true,
		},
		{
			name: "negative max element size",
			req: QueryRequest{
				SeriesIDs:    []common.SeriesID{1},
				MaxBatchSize: -1,
			},
			expectErr: true,
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			resultsCh, errCh := sidx.StreamingQuery(ctx, tt.req)
			// Drain resultsCh to wait for query to close
			for range resultsCh {
				_ = 0
			}
			err, ok := <-errCh
			if tt.expectErr {
				if ok {
					assert.Error(t, err)
				} else {
					// Validation error, should have been sent to errCh
					t.Fatal("Expected error but got none")
				}
			} else {
				if ok {
					assert.NoError(t, err)
				}
			}
		})
	}
}

func TestSIDX_StreamingQuery_MatchesBlockingQuery(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	ctx := context.Background()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 50, "series1-data50", createTestTag("env", "prod")),
		createTestWriteRequest(1, 150, "series1-data150", createTestTag("env", "prod")),
		createTestWriteRequest(1, 250, "series1-data250", createTestTag("env", "prod")),
		createTestWriteRequest(2, 75, "series2-data75", createTestTag("env", "staging")),
		createTestWriteRequest(2, 175, "series2-data175", createTestTag("env", "staging")),
		createTestWriteRequest(2, 275, "series2-data275", createTestTag("env", "staging")),
	}
	writeTestData(t, sidx, reqs, 7, 7)
	waitForIntroducerLoop()

	testCases := []struct {
		name string
		req  QueryRequest
	}{
		{
			name: "ascending_all_series",
			req: QueryRequest{
				SeriesIDs:    []common.SeriesID{1, 2},
				Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
				MaxBatchSize: 2,
				MinKey: func() *int64 {
					v := int64(50)
					return &v
				}(),
				MaxKey: func() *int64 {
					v := int64(275)
					return &v
				}(),
			},
		},
		{
			name: "descending_single_series",
			req: QueryRequest{
				SeriesIDs:    []common.SeriesID{1},
				Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_DESC},
				MaxBatchSize: 2,
			},
		},
		{
			name: "range_filtered_series",
			req: QueryRequest{
				SeriesIDs:    []common.SeriesID{2},
				Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
				MaxBatchSize: 2,
				MinKey: func() *int64 {
					v := int64(100)
					return &v
				}(),
				MaxKey: func() *int64 {
					v := int64(200)
					return &v
				}(),
			},
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			// Collect expected results from streaming query
			var (
				expectedKeys []int64
				expectedData [][]byte
				expectedSIDs []common.SeriesID
			)

			resultsCh, errCh := sidx.StreamingQuery(ctx, tc.req)
			for res := range resultsCh {
				require.NoError(t, res.Error)
				expectedKeys = append(expectedKeys, res.Keys...)
				expectedData = append(expectedData, res.Data...)
				expectedSIDs = append(expectedSIDs, res.SIDs...)
			}
			if err, ok := <-errCh; ok {
				require.NoError(t, err)
			}

			// Run the streaming query again and verify consistency
			var (
				gotKeys []int64
				gotData [][]byte
				gotSIDs []common.SeriesID
			)

			resultsCh2, errCh2 := sidx.StreamingQuery(ctx, tc.req)
			for res := range resultsCh2 {
				require.NoError(t, res.Error)
				gotKeys = append(gotKeys, res.Keys...)
				gotData = append(gotData, res.Data...)
				gotSIDs = append(gotSIDs, res.SIDs...)
			}
			if err, ok := <-errCh2; ok {
				require.NoError(t, err)
			}

			require.Equal(t, expectedKeys, gotKeys)
			require.Equal(t, expectedData, gotData)
			require.Equal(t, expectedSIDs, gotSIDs)
			require.Equal(t, len(expectedKeys), len(gotKeys))
		})
	}
}

func TestSIDX_StreamingQuery_BatchSizingAndCapacity(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 10, "trace1"),
		createTestWriteRequest(1, 11, "trace2"),
		createTestWriteRequest(1, 12, "trace3"),
		createTestWriteRequest(1, 13, "trace4"),
		createTestWriteRequest(1, 14, "trace1"),
	}
	writeTestData(t, sidx, reqs, 8, 8)
	waitForIntroducerLoop()

	queryReq := QueryRequest{
		SeriesIDs:    []common.SeriesID{1},
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 0,
	}

	resultsCh, errCh := sidx.StreamingQuery(context.Background(), queryReq)
	require.Equal(t, queryReq.MaxBatchSize, cap(resultsCh))

	var batches []*QueryResponse
	for res := range resultsCh {
		require.NoError(t, res.Error)
		batches = append(batches, res)
	}

	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	require.NotEmpty(t, batches)

	// Get expected count by collecting all streaming results
	expectedCount := 0
	resultsCh2, errCh2 := sidx.StreamingQuery(context.Background(), queryReq)
	for res := range resultsCh2 {
		require.NoError(t, res.Error)
		expectedCount += res.Len()
	}
	if err, ok := <-errCh2; ok {
		require.NoError(t, err)
	}

	totalResults := 0

	for i, batch := range batches {
		totalResults += batch.Len()

		if i < len(batches)-1 {
			require.LessOrEqual(t, batch.Len(), queryReq.MaxBatchSize)
		}

		uniqueData := make(map[string]struct{})
		for _, data := range batch.Data {
			uniqueData[string(data)] = struct{}{}
		}
		if i < len(batches)-1 {
			require.LessOrEqual(t, len(uniqueData), queryReq.MaxBatchSize)
		}
	}

	require.Equal(t, expectedCount, totalResults)
}

func TestSIDX_StreamingQuery_ChannelLifecycle(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 20, "data20"),
		createTestWriteRequest(1, 30, "data30"),
		createTestWriteRequest(1, 40, "data40"),
	}
	writeTestData(t, sidx, reqs, 9, 9)
	waitForIntroducerLoop()

	queryReq := QueryRequest{
		SeriesIDs:    []common.SeriesID{1},
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 2,
	}

	// Get expected count from first streaming query
	expectedCount := 0
	resultsCh0, errCh0 := sidx.StreamingQuery(context.Background(), queryReq)
	for res := range resultsCh0 {
		require.NoError(t, res.Error)
		expectedCount += res.Len()
	}
	if err, ok := <-errCh0; ok {
		require.NoError(t, err)
	}

	resultsCh, errCh := sidx.StreamingQuery(context.Background(), queryReq)

	totalResults := 0
	for res := range resultsCh {
		require.NoError(t, res.Error)
		totalResults += res.Len()
	}
	require.Equal(t, expectedCount, totalResults)

	_, ok := <-resultsCh
	require.False(t, ok)

	select {
	case err, ok := <-errCh:
		require.False(t, ok)
		require.NoError(t, err)
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout waiting for error channel to close")
	}
}

func TestSIDX_StreamingQuery_ErrorPropagation(t *testing.T) {
	t.Run("validation_error", func(t *testing.T) {
		sidx := createTestSIDX(t)
		defer func() {
			assert.NoError(t, sidx.Close())
		}()

		resultsCh, errCh := sidx.StreamingQuery(context.Background(), QueryRequest{})

		select {
		case err, ok := <-errCh:
			require.True(t, ok)
			require.Error(t, err)
			require.Contains(t, err.Error(), "SeriesID")
		case <-time.After(500 * time.Millisecond):
			t.Fatal("expected validation error but errCh did not emit")
		}

		// Drain resultsCh to wait for query to close
		for range resultsCh {
			_ = 0
		}
		_, ok := <-resultsCh
		require.False(t, ok)
	})

	t.Run("block_scanner_error", func(t *testing.T) {
		sidx := createTestSIDXWithOptions(t, func(opts *Options) {
			opts.Memory = &strictQuotaProtector{}
		})
		defer func() {
			assert.NoError(t, sidx.Close())
		}()

		largePayload := strings.Repeat("quota-trip", 16)
		reqs := []WriteRequest{
			createTestWriteRequest(1, 100, largePayload),
			createTestWriteRequest(1, 200, largePayload),
			createTestWriteRequest(1, 300, largePayload),
		}
		writeTestData(t, sidx, reqs, 10, 10)
		waitForIntroducerLoop()

		resultsCh, errCh := sidx.StreamingQuery(context.Background(), QueryRequest{
			SeriesIDs:    []common.SeriesID{1},
			Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			MaxBatchSize: 0,
		})

		// Drain resultsCh to wait for query to close
		for range resultsCh {
			t.Fatal("expected no results for quota exceeded")
		}

		select {
		case err, ok := <-errCh:
			require.True(t, ok)
			require.Error(t, err)
			require.Contains(t, err.Error(), "quota exceeded")
		case <-time.After(time.Second):
			t.Fatal("expected block scanner error but errCh did not emit")
		}
	})
}

func TestSIDX_StreamingQuery_EdgeCases(t *testing.T) {
	t.Run("no_results", func(t *testing.T) {
		sidx := createTestSIDX(t)
		defer func() {
			assert.NoError(t, sidx.Close())
		}()

		reqs := []WriteRequest{
			createTestWriteRequest(1, 10, "data10"),
			createTestWriteRequest(1, 20, "data20"),
		}
		writeTestData(t, sidx, reqs, 11, 11)
		waitForIntroducerLoop()

		resultsCh, errCh := sidx.StreamingQuery(context.Background(), QueryRequest{
			SeriesIDs:    []common.SeriesID{999},
			Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			MaxBatchSize: 2,
		})

		// Drain resultsCh to wait for query to close
		for range resultsCh {
			t.Fatal("expected no results for missing series")
		}

		err, ok := <-errCh
		require.False(t, ok)
		require.NoError(t, err)
	})

	t.Run("single_batch_smaller_than_limit", func(t *testing.T) {
		sidx := createTestSIDX(t)
		defer func() {
			assert.NoError(t, sidx.Close())
		}()

		reqs := []WriteRequest{
			createTestWriteRequest(2, 10, "single10"),
			createTestWriteRequest(2, 20, "single20"),
		}
		writeTestData(t, sidx, reqs, 12, 12)
		waitForIntroducerLoop()

		resultsCh, errCh := sidx.StreamingQuery(context.Background(), QueryRequest{
			SeriesIDs:    []common.SeriesID{2},
			Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			MaxBatchSize: 10,
		})

		var batches []*QueryResponse
		for res := range resultsCh {
			require.NoError(t, res.Error)
			batches = append(batches, res)
		}

		if err, ok := <-errCh; ok {
			require.NoError(t, err)
		}

		require.Len(t, batches, 1)
		require.Equal(t, len(reqs), batches[0].Len())
	})

	t.Run("mixed_memory_and_disk_parts", func(t *testing.T) {
		idx := createTestSIDX(t)
		defer func() {
			assert.NoError(t, idx.Close())
		}()

		const (
			memPartID  = uint64(100)
			diskPartID = uint64(101)
		)

		memReqs := []WriteRequest{
			createTestWriteRequest(1, 5, "mem-5"),
			createTestWriteRequest(1, 15, "mem-15"),
		}
		writeTestData(t, idx, memReqs, 13, memPartID)

		diskReqs := []WriteRequest{
			createTestWriteRequest(2, 100, "disk-100"),
			createTestWriteRequest(2, 200, "disk-200"),
		}
		writeTestData(t, idx, diskReqs, 14, diskPartID)
		waitForIntroducerLoop()

		raw := idx.(*sidx)
		flushIntro, err := raw.Flush(map[uint64]struct{}{diskPartID: {}})
		require.NoError(t, err)
		require.NotNil(t, flushIntro)
		raw.IntroduceFlushed(flushIntro)
		flushIntro.Release()

		resultsCh, errCh := idx.StreamingQuery(context.Background(), QueryRequest{
			SeriesIDs:    []common.SeriesID{1, 2},
			Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
			MaxBatchSize: 0,
		})

		var (
			collectedKeys []int64
			collectedData []string
		)

		for res := range resultsCh {
			require.NoError(t, res.Error)
			collectedKeys = append(collectedKeys, res.Keys...)
			for _, data := range res.Data {
				collectedData = append(collectedData, string(data))
			}
		}

		if err, ok := <-errCh; ok {
			require.NoError(t, err)
		}

		expectedKeys := []int64{5, 15, 100, 200}
		require.ElementsMatch(t, expectedKeys, collectedKeys)
		require.ElementsMatch(t, []string{"mem-5", "mem-15", "disk-100", "disk-200"}, collectedData)
	})
}

func TestSIDX_StreamingQuery_BatchSizeAndOrder(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 10, "data10"),
		createTestWriteRequest(1, 20, "data20"),
		createTestWriteRequest(1, 30, "data30"),
		createTestWriteRequest(1, 40, "data40"),
	}
	writeTestData(t, sidx, reqs, 42, 42)
	waitForIntroducerLoop()

	queryReq := QueryRequest{
		SeriesIDs:    []common.SeriesID{1},
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 2,
	}

	resultsCh, errCh := sidx.StreamingQuery(context.Background(), queryReq)

	var (
		allKeys    []int64
		lastKey    int64
		lastKeySet bool
		batchCount int
	)

	for res := range resultsCh {
		require.NoError(t, res.Error)
		require.NotEmpty(t, res.Keys)
		if queryReq.MaxBatchSize > 0 {
			require.LessOrEqual(t, res.Len(), queryReq.MaxBatchSize)
		}
		batchCount++

		for _, key := range res.Keys {
			if lastKeySet {
				require.GreaterOrEqual(t, key, lastKey)
			}
			lastKey = key
			lastKeySet = true
			allKeys = append(allKeys, key)
		}
	}

	if err, ok := <-errCh; ok {
		require.NoError(t, err)
	}

	require.Equal(t, len(reqs), len(allKeys))
	require.NotZero(t, batchCount)
}

func TestSIDX_StreamingQuery_ContextCancellation(t *testing.T) {
	sidx := createTestSIDX(t)
	defer func() {
		assert.NoError(t, sidx.Close())
	}()

	reqs := []WriteRequest{
		createTestWriteRequest(1, 10, "data10"),
		createTestWriteRequest(1, 20, "data20"),
		createTestWriteRequest(1, 30, "data30"),
		createTestWriteRequest(1, 40, "data40"),
		createTestWriteRequest(1, 50, "data50"),
	}
	writeTestData(t, sidx, reqs, 43, 43)
	waitForIntroducerLoop()

	ctx, cancel := context.WithCancel(context.Background())

	queryReq := QueryRequest{
		SeriesIDs:    []common.SeriesID{1},
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 1,
	}

	resultsCh, errCh := sidx.StreamingQuery(ctx, queryReq)

	done := make(chan struct{})
	go func() {
		// Drain resultsCh to wait for query to close
		for range resultsCh {
			_ = 0
		}
		close(done)
	}()

	cancel()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for streaming query to halt after cancellation")
	}

	// Depending on timing, the stream may finish before cancellation takes effect
	// (especially on faster CI runners). In that case, errCh will be closed
	// without emitting an error, which is acceptable.
	if err, ok := <-errCh; ok {
		require.ErrorIs(t, err, context.Canceled)
	}
}
