// 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 trace

import (
	"context"
	"errors"
	"strconv"
	"testing"

	"github.com/google/go-cmp/cmp"

	modelv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/model/v1"
	"github.com/apache/skywalking-banyandb/banyand/internal/sidx"
	"github.com/apache/skywalking-banyandb/banyand/queue"
	"github.com/apache/skywalking-banyandb/pkg/index"
)

type fakeSIDX struct {
	responses []*sidx.QueryResponse
}

func (f *fakeSIDX) StreamingQuery(ctx context.Context, _ sidx.QueryRequest) (<-chan *sidx.QueryResponse, <-chan error) {
	results := make(chan *sidx.QueryResponse, len(f.responses))
	errCh := make(chan error, 1)

	go func() {
		defer close(results)
		defer close(errCh)

		for _, resp := range f.responses {
			select {
			case <-ctx.Done():
				errCh <- ctx.Err()
				return
			case results <- resp:
			}
		}
	}()

	return results, errCh
}

func (f *fakeSIDX) IntroduceMemPart(uint64, *sidx.MemPart)          { panic("not implemented") }
func (f *fakeSIDX) IntroduceFlushed(*sidx.FlusherIntroduction)      {}
func (f *fakeSIDX) IntroduceMerged(*sidx.MergerIntroduction) func() { return func() {} }
func (f *fakeSIDX) ConvertToMemPart([]sidx.WriteRequest, int64, *int64, *int64) (*sidx.MemPart, error) {
	panic("not implemented")
}

func (f *fakeSIDX) Query(context.Context, sidx.QueryRequest) (*sidx.QueryResponse, error) {
	panic("not implemented")
}

func (f *fakeSIDX) QuerySync(_ context.Context, _ sidx.QueryRequest) ([]*sidx.QueryResponse, error) {
	return f.responses, nil
}
func (f *fakeSIDX) Stats(context.Context) (*sidx.Stats, error) { return &sidx.Stats{}, nil }
func (f *fakeSIDX) Close() error                               { return nil }
func (f *fakeSIDX) Flush(map[uint64]struct{}) (*sidx.FlusherIntroduction, error) {
	panic("not implemented")
}

func (f *fakeSIDX) Merge(<-chan struct{}, map[uint64]struct{}, uint64, func([]byte) bool) (*sidx.MergerIntroduction, error) {
	panic("not implemented")
}

func (f *fakeSIDX) StreamingParts(map[uint64]struct{}, string, uint32, string) ([]queue.StreamingPartData, []func()) {
	panic("not implemented")
}
func (f *fakeSIDX) PartPaths(map[uint64]struct{}) map[uint64]string { return map[uint64]string{} }
func (f *fakeSIDX) IntroduceSynced(map[uint64]struct{}) func()      { return func() {} }
func (f *fakeSIDX) TakeFileSnapshot(_ string) error                 { return nil }
func (f *fakeSIDX) ScanQuery(context.Context, sidx.ScanQueryRequest) ([]*sidx.QueryResponse, error) {
	return nil, nil
}

func (f *fakeSIDX) PrepareMemPart(uint64, *sidx.MemPart) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDX) PrepareFilePart(uint64, string) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDX) PrepareFlushed(*sidx.FlusherIntroduction) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDX) PrepareMerged(*sidx.MergerIntroduction) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDX) PrepareSynced(map[uint64]struct{}) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}
func (f *fakeSIDX) CurrentSnapshot() *sidx.Snapshot { return nil }
func (f *fakeSIDX) ReplaceSnapshot(*sidx.Snapshot)  {}

type fakeSIDXWithErr struct {
	*fakeSIDX
	err error
}

func (f *fakeSIDXWithErr) StreamingQuery(ctx context.Context, _ sidx.QueryRequest) (<-chan *sidx.QueryResponse, <-chan error) {
	results := make(chan *sidx.QueryResponse, len(f.responses))
	errCh := make(chan error, 1)

	go func() {
		defer close(results)
		defer close(errCh)

		for _, resp := range f.responses {
			select {
			case <-ctx.Done():
				return
			case results <- resp:
			}
		}

		// Send error after all results have been sent
		// This tests that errors are propagated even after data processing
		if f.err != nil {
			select {
			case <-ctx.Done():
			case errCh <- f.err:
			}
		}
	}()

	return results, errCh
}

func encodeTraceIDForTest(id string) []byte {
	buf := make([]byte, len(id)+1)
	buf[0] = byte(idFormatV1)
	copy(buf[1:], id)
	return buf
}

func TestStreamSIDXTraceBatches_ProducesOrderedBatches(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 2,
	}

	responses := []*sidx.QueryResponse{
		{
			Keys: []int64{1, 2},
			Data: [][]byte{
				encodeTraceIDForTest("a"),
				encodeTraceIDForTest("b"),
			},
			PartIDs: []uint64{100, 101},
		},
		{
			Keys: []int64{2, 3},
			Data: [][]byte{
				encodeTraceIDForTest("b"),
				encodeTraceIDForTest("c"),
			},
			PartIDs: []uint64{101, 102},
		},
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{&fakeSIDX{responses: responses}}, req, 3)

	var batches []traceBatch
	for batch := range batchCh {
		if batch.err != nil {
			t.Fatalf("unexpected error batch: %v", batch.err)
		}
		batches = append(batches, batch)
	}

	if len(batches) != 2 {
		t.Fatalf("expected 2 batches, got %d", len(batches))
	}

	wantBatch0 := map[uint64][]string{
		100: {"a"},
		101: {"b"},
	}
	if diff := cmp.Diff(wantBatch0, batches[0].traceIDs); diff != "" {
		t.Fatalf("first batch mismatch (-want +got):\n%s", diff)
	}

	wantBatch1 := map[uint64][]string{
		102: {"c"},
	}
	if diff := cmp.Diff(wantBatch1, batches[1].traceIDs); diff != "" {
		t.Fatalf("second batch mismatch (-want +got):\n%s", diff)
	}

	wantKeys := map[string]int64{"a": 1, "b": 2, "c": 3}
	for _, batch := range batches {
		for _, ids := range batch.traceIDs {
			for _, tid := range ids {
				if got := batch.keys[tid]; got != wantKeys[tid] {
					t.Fatalf("unexpected key for %s: got %d, want %d", tid, got, wantKeys[tid])
				}
			}
		}
	}
}

func TestStreamSIDXTraceBatches_OrdersDescending(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_DESC},
		MaxBatchSize: 2,
	}

	responses := []*sidx.QueryResponse{
		{
			Keys: []int64{5, 4},
			Data: [][]byte{
				encodeTraceIDForTest("e"),
				encodeTraceIDForTest("d"),
			},
			PartIDs: []uint64{200, 201},
		},
		{
			Keys: []int64{4, 3},
			Data: [][]byte{
				encodeTraceIDForTest("d"),
				encodeTraceIDForTest("c"),
			},
			PartIDs: []uint64{201, 202},
		},
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{&fakeSIDX{responses: responses}}, req, 3)

	var batches []traceBatch
	for batch := range batchCh {
		if batch.err != nil {
			t.Fatalf("unexpected error batch: %v", batch.err)
		}
		batches = append(batches, batch)
	}

	if len(batches) != 2 {
		t.Fatalf("expected 2 batches, got %d", len(batches))
	}

	wantBatch0 := map[uint64][]string{
		200: {"e"},
		201: {"d"},
	}
	if diff := cmp.Diff(wantBatch0, batches[0].traceIDs); diff != "" {
		t.Fatalf("first batch mismatch (-want +got):\n%s", diff)
	}

	wantBatch1 := map[uint64][]string{
		202: {"c"},
	}
	if diff := cmp.Diff(wantBatch1, batches[1].traceIDs); diff != "" {
		t.Fatalf("second batch mismatch (-want +got):\n%s", diff)
	}

	wantKeys := map[string]int64{"e": 5, "d": 4, "c": 3}
	for _, batch := range batches {
		for _, ids := range batch.traceIDs {
			for _, tid := range ids {
				if got := batch.keys[tid]; got != wantKeys[tid] {
					t.Fatalf("unexpected key for %s: got %d, want %d", tid, got, wantKeys[tid])
				}
			}
		}
	}
}

func TestStreamSIDXTraceBatches_PropagatesErrorAfterCancellation(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 1,
	}

	streamErr := errors.New("stream failure")

	sidxInstance := &fakeSIDXWithErr{
		fakeSIDX: &fakeSIDX{
			responses: []*sidx.QueryResponse{
				{
					Keys: []int64{1},
					Data: [][]byte{
						encodeTraceIDForTest("trace-1"),
					},
					PartIDs: []uint64{300},
				},
			},
		},
		err: streamErr,
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0)

	var (
		dataSeen bool
		errBatch traceBatch
		errSeen  bool
	)

	for batch := range batchCh {
		if batch.err != nil {
			errBatch = batch
			errSeen = true
			continue
		}
		for _, ids := range batch.traceIDs {
			if len(ids) > 0 {
				dataSeen = true
				break
			}
		}
	}

	// The error must be propagated
	if !errSeen {
		t.Fatalf("expected error batch but none received")
	}
	if !errors.Is(errBatch.err, streamErr) {
		t.Fatalf("unexpected error: %v", errBatch.err)
	}

	// Data may or may not be seen depending on error timing.
	// If the error arrives on errEvents before the first loop iteration processes
	// the heap, dataSeen will be false. Both outcomes are valid.
	// The important invariant is that the error is always propagated.
	if dataSeen {
		t.Logf("data was processed before error detected (typical case)")
	} else {
		t.Logf("error detected before data processing (race condition)")
	}
}

// fakeSIDXWithImmediateError simulates errors from blockScanResultBatch.
// by sending errors immediately via errCh (like scan errors would).
type fakeSIDXWithImmediateError struct {
	*fakeSIDX
	err error
}

func (f *fakeSIDXWithImmediateError) StreamingQuery(ctx context.Context, _ sidx.QueryRequest) (<-chan *sidx.QueryResponse, <-chan error) {
	results := make(chan *sidx.QueryResponse, len(f.responses))
	errCh := make(chan error, 1)

	go func() {
		defer close(results)
		defer close(errCh)

		// Send error immediately (simulating blockScanner.scan error)
		if f.err != nil {
			errCh <- f.err
		}

		// Then send any responses
		for _, resp := range f.responses {
			select {
			case <-ctx.Done():
				return
			case results <- resp:
			}
		}
	}()

	return results, errCh
}

// TestStreamSIDXTraceBatches_PropagatesBlockScannerError verifies that errors.
// from blockScanResultBatch.err (sent via errCh) are properly propagated through
// the streaming pipeline to the traceBatch channel and ultimately to queryResult.Pull().
func TestStreamSIDXTraceBatches_PropagatesBlockScannerError(t *testing.T) {
	tests := []struct {
		name         string
		scanError    string
		errorContain string
	}{
		{
			name:         "iterator_init_error",
			scanError:    "cannot init iter: iterator initialization failed",
			errorContain: "cannot init iter",
		},
		{
			name:         "block_scan_error",
			scanError:    "sidx block scan quota exceeded: used 1000 bytes, quota is 500 bytes",
			errorContain: "quota exceeded",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := sidx.QueryRequest{
				Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
				MaxBatchSize: 2,
			}

			scanErr := errors.New(tt.scanError)

			// Use fakeSIDXWithImmediateError - since errCh and results channel
			// race each other, we don't test ordering, just that error propagates
			sidxInstance := &fakeSIDXWithImmediateError{
				fakeSIDX: &fakeSIDX{
					responses: nil,
				},
				err: scanErr,
			}

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

			var tr trace
			batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0)

			var receivedError error

			// Consume all batches until we receive the error
			for batch := range batchCh {
				if batch.err != nil {
					receivedError = batch.err
					// Don't break - consume remaining batches to avoid goroutine leaks
					continue
				}
			}

			// Verify error was eventually propagated
			if receivedError == nil {
				t.Fatalf("expected error containing %q but got none", tt.errorContain)
			}
			if !errors.Is(receivedError, scanErr) {
				// Check if error is wrapped
				errMsg := receivedError.Error()
				if errMsg == "" || !contains(errMsg, tt.errorContain) {
					t.Fatalf("expected error containing %q but got: %v", tt.errorContain, receivedError)
				}
			}
		})
	}
}

// TestStreamSIDXTraceBatches_DrainErrorEventsGuaranteed verifies that.
// drainErrorEvents is always called via defer, even on early returns.
func TestStreamSIDXTraceBatches_DrainErrorEventsGuaranteed(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 10,
	}

	scanErr := errors.New("scan error from defer test")

	// Create fake that will send error after some results
	sidxInstance := &fakeSIDXWithImmediateError{
		fakeSIDX: &fakeSIDX{
			responses: []*sidx.QueryResponse{
				{
					Keys: []int64{1},
					Data: [][]byte{
						encodeTraceIDForTest("trace-1"),
					},
					PartIDs: []uint64{400},
				},
			},
		},
		err: scanErr,
	}

	ctx, cancel := context.WithCancel(context.Background())
	// Cancel context immediately to force early return
	cancel()

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0)

	var receivedError error

	// Even with canceled context, we should receive the error
	// because defer ensures drainErrorEvents is called
	for batch := range batchCh {
		if batch.err != nil {
			receivedError = batch.err
		}
	}

	// The error might not be propagated if context is canceled immediately
	// but this test verifies the defer mechanism exists
	t.Logf("Received error (may be nil due to immediate cancellation): %v", receivedError)
}

// TestStreamSIDXTraceBatches_ErrorEmissionResilience verifies that.
// emitError tries to send even when context is canceled.
func TestStreamSIDXTraceBatches_ErrorEmissionResilience(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 1,
	}

	scanErr := errors.New("scan error during processing")

	sidxInstance := &fakeSIDXWithImmediateError{
		fakeSIDX: &fakeSIDX{
			responses: []*sidx.QueryResponse{
				{
					Keys: []int64{1, 2, 3},
					Data: [][]byte{
						encodeTraceIDForTest("trace-1"),
						encodeTraceIDForTest("trace-2"),
						encodeTraceIDForTest("trace-3"),
					},
					PartIDs: []uint64{500, 501, 502},
				},
			},
		},
		err: scanErr,
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0)

	var receivedError error
	batchCount := 0

	// Read first batch, then wait a bit for error to be sent to errEvents
	for batch := range batchCh {
		if batch.err != nil {
			receivedError = batch.err
			break
		}
		batchCount++
		// After first batch, continue reading to trigger drainErrorEvents
	}

	// We should eventually get the error either from pollErrEvents or drainErrorEvents
	if receivedError == nil {
		t.Fatal("expected error to be propagated but got none")
	}

	if !errors.Is(receivedError, scanErr) && !contains(receivedError.Error(), "scan error") {
		t.Fatalf("unexpected error: %v", receivedError)
	}
}

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsHelper(s, substr))
}

func containsHelper(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}

// fakeSIDXInfinite simulates a SIDX that returns an infinite stream of results.
// It continues to generate responses until the context is canceled.
type fakeSIDXInfinite struct {
	traceIDPrefix string
	batchSize     int
	keyStart      int64
}

func (f *fakeSIDXInfinite) ScanQuery(context.Context, sidx.ScanQueryRequest) ([]*sidx.QueryResponse, error) {
	return nil, nil
}

func (f *fakeSIDXInfinite) QuerySync(_ context.Context, _ sidx.QueryRequest) ([]*sidx.QueryResponse, error) {
	return nil, nil
}

func (f *fakeSIDXInfinite) StreamingQuery(ctx context.Context, _ sidx.QueryRequest) (<-chan *sidx.QueryResponse, <-chan error) {
	results := make(chan *sidx.QueryResponse)
	errCh := make(chan error, 1)

	go func() {
		defer close(results)
		defer close(errCh)

		key := f.keyStart
		counter := 0

		for {
			select {
			case <-ctx.Done():
				errCh <- ctx.Err()
				return
			default:
				batchSize := f.batchSize
				if batchSize <= 0 {
					batchSize = 10
				}

				keys := make([]int64, batchSize)
				data := make([][]byte, batchSize)
				partIDs := make([]uint64, batchSize)

				for i := 0; i < batchSize; i++ {
					keys[i] = key
					prefix := f.traceIDPrefix
					if prefix == "" {
						prefix = "trace"
					}
					traceID := prefix + "-" + strconv.Itoa(counter)
					data[i] = encodeTraceIDForTest(traceID)
					partIDs[i] = uint64(1000 + counter)
					key++
					counter++
				}

				resp := &sidx.QueryResponse{
					Keys:    keys,
					Data:    data,
					PartIDs: partIDs,
				}

				select {
				case <-ctx.Done():
					errCh <- ctx.Err()
					return
				case results <- resp:
				}
			}
		}
	}()

	return results, errCh
}

func (f *fakeSIDXInfinite) IntroduceMemPart(uint64, *sidx.MemPart)          { panic("not implemented") }
func (f *fakeSIDXInfinite) IntroduceFlushed(*sidx.FlusherIntroduction)      {}
func (f *fakeSIDXInfinite) IntroduceMerged(*sidx.MergerIntroduction) func() { return func() {} }
func (f *fakeSIDXInfinite) ConvertToMemPart([]sidx.WriteRequest, int64, *int64, *int64) (*sidx.MemPart, error) {
	panic("not implemented")
}

func (f *fakeSIDXInfinite) Query(context.Context, sidx.QueryRequest) (*sidx.QueryResponse, error) {
	panic("not implemented")
}
func (f *fakeSIDXInfinite) Stats(context.Context) (*sidx.Stats, error) { return &sidx.Stats{}, nil }
func (f *fakeSIDXInfinite) Close() error                               { return nil }
func (f *fakeSIDXInfinite) Flush(map[uint64]struct{}) (*sidx.FlusherIntroduction, error) {
	panic("not implemented")
}

func (f *fakeSIDXInfinite) Merge(<-chan struct{}, map[uint64]struct{}, uint64, func([]byte) bool) (*sidx.MergerIntroduction, error) {
	panic("not implemented")
}

func (f *fakeSIDXInfinite) StreamingParts(map[uint64]struct{}, string, uint32, string) ([]queue.StreamingPartData, []func()) {
	panic("not implemented")
}

func (f *fakeSIDXInfinite) PartPaths(map[uint64]struct{}) map[uint64]string {
	return map[uint64]string{}
}
func (f *fakeSIDXInfinite) IntroduceSynced(map[uint64]struct{}) func() { return func() {} }
func (f *fakeSIDXInfinite) TakeFileSnapshot(_ string) error            { return nil }
func (f *fakeSIDXInfinite) PrepareMemPart(uint64, *sidx.MemPart) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDXInfinite) PrepareFilePart(uint64, string) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDXInfinite) PrepareFlushed(*sidx.FlusherIntroduction) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDXInfinite) PrepareMerged(*sidx.MergerIntroduction) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}

func (f *fakeSIDXInfinite) PrepareSynced(map[uint64]struct{}) func(cur *sidx.Snapshot) *sidx.Snapshot {
	return func(cur *sidx.Snapshot) *sidx.Snapshot { return cur }
}
func (f *fakeSIDXInfinite) CurrentSnapshot() *sidx.Snapshot { return nil }
func (f *fakeSIDXInfinite) ReplaceSnapshot(*sidx.Snapshot)  {}

// TestStreamSIDXTraceBatches_InfiniteChannelContinuesUntilCanceled verifies that
// the streaming pipeline continues streaming from an infinite channel until context is canceled.
func TestStreamSIDXTraceBatches_InfiniteChannelContinuesUntilCanceled(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 5,
	}

	sidxInstance := &fakeSIDXInfinite{
		batchSize:     10,
		keyStart:      1,
		traceIDPrefix: "inf",
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0 /* no maxTraceSize limit */)

	const targetTraceIDs = 50
	totalTraceIDs := 0
	seenIDs := make(map[string]struct{})

	for batch := range batchCh {
		if batch.err != nil {
			// Context cancellation is expected
			if errors.Is(batch.err, context.Canceled) {
				break
			}
			t.Fatalf("unexpected error batch: %v", batch.err)
		}

		for _, ids := range batch.traceIDs {
			for _, tid := range ids {
				if _, exists := seenIDs[tid]; exists {
					t.Fatalf("duplicate trace ID: %s", tid)
				}
				seenIDs[tid] = struct{}{}
				totalTraceIDs++
			}
		}

		// Cancel after we've received enough traces to prove it's streaming
		if totalTraceIDs >= targetTraceIDs {
			cancel()
		}
	}

	// Should have received at least targetTraceIDs results
	if totalTraceIDs < targetTraceIDs {
		t.Fatalf("expected at least %d trace IDs, got %d", targetTraceIDs, totalTraceIDs)
	}

	t.Logf("Successfully received %d unique trace IDs from infinite stream before cancellation", totalTraceIDs)
}

// TestStreamSIDXTraceBatches_InfiniteChannelCancellation verifies that
// canceling the context properly stops all goroutines in the streaming pipeline
// and cancels the SIDX context to close the infinite channel.
func TestStreamSIDXTraceBatches_InfiniteChannelCancellation(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 5,
	}

	sidxInstance := &fakeSIDXInfinite{
		batchSize:     10,
		keyStart:      1,
		traceIDPrefix: "cancel",
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0 /* no maxTraceSize limit */)

	// Read a few batches, then cancel
	batchesRead := 0
	const batchesToRead = 3

	for batch := range batchCh {
		if batch.err != nil {
			// Context cancellation error is expected
			if errors.Is(batch.err, context.Canceled) {
				t.Logf("Received expected cancellation error: %v", batch.err)
				break
			}
			t.Fatalf("unexpected error: %v", batch.err)
		}

		batchesRead++
		if batchesRead >= batchesToRead {
			// Cancel the context to stop the infinite stream
			cancel()
		}
	}

	// Verify we read some batches before cancellation
	if batchesRead < batchesToRead {
		t.Fatalf("expected to read at least %d batches, got %d", batchesToRead, batchesRead)
	}

	t.Logf("Successfully canceled infinite stream after reading %d batches", batchesRead)
}

// TestStreamSIDXTraceBatches_InfiniteChannelGoroutineCleanup verifies that
// all goroutines are properly cleaned up when context is canceled.
func TestStreamSIDXTraceBatches_InfiniteChannelGoroutineCleanup(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 5,
	}

	sidxInstance := &fakeSIDXInfinite{
		batchSize:     5,
		keyStart:      1,
		traceIDPrefix: "cleanup",
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, []sidx.SIDX{sidxInstance}, req, 0 /* no limit */)

	totalTraceIDs := 0
	batchesRead := 0

	for batch := range batchCh {
		if batch.err != nil {
			// Context cancellation is expected
			if errors.Is(batch.err, context.Canceled) {
				t.Logf("Received expected cancellation error")
				break
			}
			t.Fatalf("unexpected error: %v", batch.err)
		}

		for _, ids := range batch.traceIDs {
			totalTraceIDs += len(ids)
		}
		batchesRead++

		// Cancel after reading a few batches
		if batchesRead >= 3 {
			cancel()
		}
	}

	// Channel should be closed
	_, ok := <-batchCh
	if ok {
		t.Fatal("channel should be closed")
	}

	if totalTraceIDs == 0 {
		t.Fatal("expected some trace IDs before cancellation")
	}

	t.Logf("Successfully cleaned up: read %d batches, %d trace IDs", batchesRead, totalTraceIDs)
}

// TestStreamSIDXTraceBatches_MultipleInfiniteSIDX verifies that
// the streaming pipeline can handle multiple infinite SIDX instances
// and properly merge their results until context is canceled.
func TestStreamSIDXTraceBatches_MultipleInfiniteSIDX(t *testing.T) {
	req := sidx.QueryRequest{
		Order:        &index.OrderBy{Sort: modelv1.Sort_SORT_ASC},
		MaxBatchSize: 10,
	}

	const targetTraceIDs = 50

	// Create multiple infinite SIDX instances with different key ranges
	sidxInstances := []sidx.SIDX{
		&fakeSIDXInfinite{
			batchSize:     5,
			keyStart:      1, // Keys: 1, 2, 3, ...
			traceIDPrefix: "s1",
		},
		&fakeSIDXInfinite{
			batchSize:     5,
			keyStart:      1000, // Keys: 1000, 1001, 1002, ...
			traceIDPrefix: "s2",
		},
		&fakeSIDXInfinite{
			batchSize:     5,
			keyStart:      2000, // Keys: 2000, 2001, 2002, ...
			traceIDPrefix: "s3",
		},
	}

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

	var tr trace
	batchCh, _ := tr.streamSIDXTraceBatches(ctx, sidxInstances, req, 0 /* no limit */)

	totalTraceIDs := 0
	seenIDs := make(map[string]struct{})
	var keys []int64

	for batch := range batchCh {
		if batch.err != nil {
			// Context cancellation is expected
			if errors.Is(batch.err, context.Canceled) {
				break
			}
			t.Fatalf("unexpected error batch: %v", batch.err)
		}

		// Use traceIDsOrder to maintain the correct order from SIDX stream
		for _, tid := range batch.traceIDsOrder {
			if _, exists := seenIDs[tid]; exists {
				t.Fatalf("duplicate trace ID: %s", tid)
			}
			seenIDs[tid] = struct{}{}
			keys = append(keys, batch.keys[tid])
			totalTraceIDs++
		}

		// Cancel after we've received enough to verify the merge is working
		if totalTraceIDs >= targetTraceIDs {
			cancel()
		}
	}

	if totalTraceIDs < targetTraceIDs {
		t.Fatalf("expected at least %d trace IDs, got %d", targetTraceIDs, totalTraceIDs)
	}

	// Verify keys are in ascending order (due to heap merge)
	for i := 1; i < len(keys); i++ {
		if keys[i] < keys[i-1] {
			t.Fatalf("keys not in ascending order: keys[%d]=%d, keys[%d]=%d", i-1, keys[i-1], i, keys[i])
		}
	}

	t.Logf("Successfully merged %d trace IDs from %d infinite SIDX instances", totalTraceIDs, len(sidxInstances))
}
