// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The 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 delegator

import (
	"context"
	"sync/atomic"
	"testing"

	"github.com/cockroachdb/errors"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
	"google.golang.org/protobuf/proto"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/querynodev2/cluster"
	"github.com/milvus-io/milvus/internal/querynodev2/segments"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/internal/util/searchutil/optimizers"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
	"github.com/milvus-io/milvus/pkg/v3/util/metric"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)

type TwoStageSearchSuite struct {
	suite.Suite

	collectionID  int64
	partitionIDs  []int64
	replicaID     int64
	vchannelName  string
	version       int64
	workerManager *cluster.MockManager
	manager       *segments.Manager
	loader        *segments.MockLoader
	chunkManager  storage.ChunkManager
	rootPath      string

	delegator *shardDelegator
}

func (s *TwoStageSearchSuite) SetupSuite() {
	paramtable.Init()
}

func (s *TwoStageSearchSuite) SetupTest() {
	s.collectionID = 1000
	s.partitionIDs = []int64{500, 501}
	s.replicaID = 65535
	s.vchannelName = "rootcoord-dml_1000_v0"
	s.version = 2000
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.rootPath = "delegator_twostage_test"

	// init schema
	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
		Name: "TestCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			},
			{
				Name:         "vector",
				FieldID:      101,
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_FloatVector,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
			},
		},
	}, &segcorepb.CollectionIndexMeta{
		MaxIndexRowCount: 100,
		IndexMetas: []*segcorepb.FieldIndexMeta{
			{
				FieldID:      101,
				CollectionID: s.collectionID,
				IndexName:    "vector_index",
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
				IndexParams: []*commonpb.KeyValuePair{
					{
						Key:   common.IndexTypeKey,
						Value: "HNSW",
					},
					{
						Key:   common.MetricTypeKey,
						Value: metric.L2,
					},
				},
			},
		},
	}, &querypb.LoadMetaInfo{
		PartitionIDs: s.partitionIDs,
	})

	// init chunkManager
	chunkManagerFactory := storage.NewTestChunkManagerFactory(paramtable.Get(), s.rootPath)
	s.chunkManager, _ = chunkManagerFactory.NewPersistentStorageChunkManager(context.Background())

	var err error
	delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
	s.Require().NoError(err)
	s.delegator = delegator.(*shardDelegator)
}

func (s *TwoStageSearchSuite) TearDownTest() {
	if s.delegator != nil {
		s.delegator.Close()
		s.delegator = nil
	}
}

func (s *TwoStageSearchSuite) TestShouldUseTwoStageSearch() {
	s.Run("disabled", func() {
		// Ensure two-stage search is disabled
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "false")
		defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       3000,
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.False(result, "should return false when two-stage search is disabled")
	})

	s.Run("segments_and_topk_below_threshold", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       1000, // below min topk
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// sealedNum=3 is below min segments (5) and topk=1000 is below min topk (2000)
		result := optimizers.ShouldUseTwoStageSearch(req, 3)
		s.False(result, "should return false when both segments and topk are below threshold")
	})

	s.Run("topk_below_threshold", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       1000, // below min topk
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// topk=1000 is below min topk (2000), so it should not pass even if sealedNum is above threshold
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.False(result, "should return false when topk is below threshold")
	})

	s.Run("segments_below_threshold", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       3000, // above min topk
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// sealedNum=3 is below min segments (5), so it should not pass even if topk is above threshold
		result := optimizers.ShouldUseTwoStageSearch(req, 3)
		s.False(result, "should return false when segments are below threshold")
	})

	s.Run("wrong_search_type_no_filter", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       3000,
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_NO_FILTER, // wrong search type
			},
		}
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.False(result, "should return false for PURE_ANN_SEARCH_NO_FILTER")
	})

	s.Run("wrong_search_type_default", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       3000,
				SearchType: internalpb.SearchType_DEFAULT, // default search type
			},
		}
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.False(result, "should return false for DEFAULT search type")
	})

	s.Run("all_conditions_met", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       3000,
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.True(result, "should return true when all conditions are met")
	})
}

// MockQueryHook implements optimizers.QueryHook for testing
type MockQueryHook struct {
	mock.Mock
}

func (m *MockQueryHook) Run(params map[string]any) error {
	args := m.Called(params)
	return args.Error(0)
}

func (m *MockQueryHook) Init(s string) error {
	args := m.Called(s)
	return args.Error(0)
}

func (m *MockQueryHook) InitTuningConfig(config map[string]string) error {
	args := m.Called(config)
	return args.Error(0)
}

func (m *MockQueryHook) DeleteTuningConfig(s string) error {
	args := m.Called(s)
	return args.Error(0)
}

func (m *MockQueryHook) CalculateEffectiveSegmentNum(rowCounts []int64, topk int64) int {
	args := m.Called(rowCounts, topk)
	return args.Int(0)
}

var _ optimizers.QueryHook = (*MockQueryHook)(nil)

func (s *TwoStageSearchSuite) TestExecuteFilterStage() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("filter_only_flag_set", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		// Verify that the request has FilterOnly=true and EnableExprCache=true
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.True(req.GetFilterOnly(), "FilterOnly should be true in filter stage")
				s.True(req.GetEnableExprCache(), "EnableExprCache should be true in filter stage")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{100},
		}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 1000,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 1000, NodeID: 1},
				},
			},
		}
		sealedRowCount := map[int64]int64{1000: 10000}

		validCounts, err := s.delegator.executeFilterStage(ctx, req, sealed, sealedRowCount)
		s.NoError(err)
		s.NotNil(validCounts)
		s.Len(validCounts, 1)
		s.Equal(int64(100), validCounts[0])
	})

	s.Run("filter_stage_error", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		mockErr := errors.New("filter stage error")
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(nil, mockErr)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 1001,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 1001, NodeID: 1},
				},
			},
		}
		sealedRowCount := map[int64]int64{1001: 10000}

		validCounts, err := s.delegator.executeFilterStage(ctx, req, sealed, sealedRowCount)
		s.Error(err)
		s.Nil(validCounts)
	})
}

func (s *TwoStageSearchSuite) TestTwoStageSearch() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("filter_stage_fails", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		mockErr := errors.New("filter stage failed")
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(nil, mockErr)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 2000,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 2000, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{2000: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.Error(err)
		s.False(fallback)
		s.Nil(results)
		s.Contains(err.Error(), "filter stage failed")
	})

	s.Run("two_stage_search_success", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		var callCount atomic.Int32
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				n := callCount.Add(1)
				if n == 1 {
					// First call should be filter stage with FilterOnly=true
					s.True(req.GetFilterOnly(), "First call should have FilterOnly=true")
					s.True(req.GetEnableExprCache(), "First call should have EnableExprCache=true")
				} else {
					// Second call should be normal search with FilterOnly=false but EnableExprCache=true
					s.False(req.GetFilterOnly(), "Second call should have FilterOnly=false")
					s.True(req.GetEnableExprCache(), "Second call should have EnableExprCache=true")
				}
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{100},
		}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 3000,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 3000, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{3000: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.False(fallback)
		s.NotNil(results)
		// Should have called SearchSegments twice: once for filter stage, once for vector search
		s.Equal(int32(2), callCount.Load(), "Should have called SearchSegments twice")
	})

	s.Run("incomplete_filter_counts_fallback_restores_enable_expr_cache", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.True(req.GetFilterOnly(), "FilterOnly should be true in filter stage")
				s.True(req.GetEnableExprCache(), "EnableExprCache should be true in filter stage")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{100},
		}, nil).Once()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		s.delegator.distribution.AddDistributions(
			SegmentEntry{NodeID: 1, SegmentID: 3100},
			SegmentEntry{NodeID: 1, SegmentID: 3101},
		)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels:     []string{s.vchannelName},
			EnableExprCache: false,
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 3100, NodeID: 1},
					{SegmentID: 3101, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{
			3100: 10000,
			3101: 10000,
		}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.True(fallback)
		s.Nil(results)
		s.False(req.GetFilterOnly(), "FilterOnly should be restored after fallback")
		s.False(req.GetEnableExprCache(), "EnableExprCache should be restored after fallback")
	})

	s.Run("optimizer_path_exercised", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.delegator.queryHook = nil
		}()

		// Enable AutoIndex so OptimizeSearchParams doesn't early-return
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
		defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)

		// Wire up a real MockQueryHook on the delegator
		mockHook := &MockQueryHook{}
		hookCalled := false
		mockHook.On("CalculateEffectiveSegmentNum", mock.Anything, mock.Anything).Return(1)
		mockHook.On("Run", mock.Anything).Run(func(args mock.Arguments) {
			hookCalled = true
			params := args.Get(0).(map[string]any)
			// In the second stage (isSecondStageSearch=true), WithFilterKey must be false
			s.False(params[common.WithFilterKey].(bool),
				"WithFilterKey should be false in second-stage search")
		}).Return(nil)
		s.delegator.queryHook = mockHook

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		var callCount atomic.Int32
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				n := callCount.Add(1)
				if n == 2 {
					optimizedPlan := &planpb.PlanNode{}
					err := proto.Unmarshal(req.GetReq().GetSerializedExprPlan(), optimizedPlan)
					s.Require().NoError(err)
					s.NotNil(optimizedPlan.GetVectorAnns().GetPredicates(),
						"second-stage search must keep predicates in the serialized plan")
				}
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{5000},
		}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 3500,
		})

		// Build a real serialized plan with VectorAnns so the optimizer doesn't skip it
		plan := &planpb.PlanNode{
			Node: &planpb.PlanNode_VectorAnns{
				VectorAnns: &planpb.VectorANNS{
					Predicates: &planpb.Expr{}, // non-nil → withFilter=true
					QueryInfo: &planpb.QueryInfo{
						Topk:           100,
						SearchParams:   `{"ef": 128}`,
						GroupByFieldId: -1,
					},
				},
			},
		}
		serializedPlan, err := proto.Marshal(plan)
		s.Require().NoError(err)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:               100,
				SerializedExprPlan: serializedPlan,
			},
			TotalChannelNum: 1,
			DmlChannels:     []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 3500, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{3500: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.False(fallback)
		s.NotNil(results)
		s.Equal(int32(2), callCount.Load(), "Should have called SearchSegments twice")
		s.True(hookCalled, "QueryHook.Run should have been called by the optimizer")
		mockHook.AssertExpectations(s.T())
	})

	s.Run("optimizer_error_propagated", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.delegator.queryHook = nil
		}()

		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
		defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)

		mockHook := &MockQueryHook{}
		mockHook.On("CalculateEffectiveSegmentNum", mock.Anything, mock.Anything).Return(1)
		mockHook.On("Run", mock.Anything).Return(errors.New("optimizer failed"))
		s.delegator.queryHook = mockHook

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		// Filter stage succeeds
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(&internalpb.SearchResults{
				FilterValidCounts: []int64{5000},
			}, nil).Once()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 3600,
		})

		plan := &planpb.PlanNode{
			Node: &planpb.PlanNode_VectorAnns{
				VectorAnns: &planpb.VectorANNS{
					QueryInfo: &planpb.QueryInfo{
						Topk:         100,
						SearchParams: `{"ef": 128}`,
					},
				},
			},
		}
		serializedPlan, err := proto.Marshal(plan)
		s.Require().NoError(err)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:               100,
				SerializedExprPlan: serializedPlan,
			},
			TotalChannelNum: 1,
			DmlChannels:     []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 3600, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{3600: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.Error(err)
		s.False(fallback)
		s.Nil(results)
		s.Contains(err.Error(), "optimizer failed")
	})

	s.Run("vector_search_stage_fails", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		mockErr := errors.New("vector search failed")
		// First call for filter stage succeeds, second call for vector search fails
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(&internalpb.SearchResults{
				FilterValidCounts: []int64{100},
			}, nil).Once()
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(nil, mockErr).Once()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 4000,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 4000, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{4000: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.Error(err)
		s.False(fallback)
		s.Nil(results)
	})
}

func (s *TwoStageSearchSuite) TestTwoStageSearchWithMultipleSegments() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("multiple_segments_filter_stats_aggregation", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}
		workers[1] = worker1
		workers[2] = worker2

		// Worker 1 handles segments on node 1
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(&internalpb.SearchResults{
				FilterValidCounts: []int64{50, 100}, // Two segments with different filter valid counts
			}, nil)

		// Worker 2 handles segments on node 2
		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(&internalpb.SearchResults{
				FilterValidCounts: []int64{200},
			}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments on multiple nodes
		s.delegator.distribution.AddDistributions(
			SegmentEntry{NodeID: 1, SegmentID: 5000},
			SegmentEntry{NodeID: 1, SegmentID: 5001},
			SegmentEntry{NodeID: 2, SegmentID: 5002},
		)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 5000, NodeID: 1},
					{SegmentID: 5001, NodeID: 1},
				},
			},
			{
				NodeID: 2,
				Segments: []SegmentEntry{
					{SegmentID: 5002, NodeID: 2},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{
			5000: 10000,
			5001: 10000,
			5002: 10000,
		}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.False(fallback)
		s.NotNil(results)
	})
}

func (s *TwoStageSearchSuite) TestShouldUseTwoStageSearchEdgeCases() {
	s.Run("topk_below_threshold_at_exact_segments", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       1999, // exactly one below threshold
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// sealedNum=5 is exactly at threshold, but topk is below threshold
		result := optimizers.ShouldUseTwoStageSearch(req, 5)
		s.False(result, "should return false when topk is below threshold")
	})

	s.Run("segments_below_threshold_at_exact_topk", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "2000")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       2000, // exactly at threshold
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// topk=2000 is exactly at threshold, but sealedNum=4 is below threshold
		result := optimizers.ShouldUseTwoStageSearch(req, 4)
		s.False(result, "should return false when segments are below threshold")
	})

	s.Run("zero_segments", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "100")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "1")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       50, // below threshold
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		result := optimizers.ShouldUseTwoStageSearch(req, 0)
		s.False(result, "should return false when both segments and topk are below threshold")
	})

	s.Run("zero_topk", func() {
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key, "true")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key, "100")
		paramtable.Get().Save(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key, "5")
		defer func() {
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchEnabled.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinTopk.Key)
			paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.TwoStageSearchMinNumSegments.Key)
		}()

		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk:       0, // zero topk
				SearchType: internalpb.SearchType_PURE_ANN_SEARCH_WITH_FILTER,
			},
		}
		// topk=0 is below threshold, so it should not pass even if sealedNum is above threshold
		result := optimizers.ShouldUseTwoStageSearch(req, 10)
		s.False(result, "should return false when topk is below threshold")
	})
}

func (s *TwoStageSearchSuite) TestTwoStageSearchWithGrowingSegments() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("growing_segments_included", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		var callCount atomic.Int32
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				callCount.Add(1)
				// All calls in two-stage search should have EnableExprCache=true
				s.True(req.GetEnableExprCache(), "EnableExprCache should be true for all two-stage search calls")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{100},
		}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(SegmentEntry{
			NodeID:    1,
			SegmentID: 6000,
		})

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 6000, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{
			{SegmentID: 6001, NodeID: 1},
		}
		sealedRowCount := map[int64]int64{6000: 10000}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.False(fallback)
		s.NotNil(results)
		// Should have called SearchSegments 3 times:
		// filter stage: 1 call for sealed only (growing not included in filter stage)
		// search stage: 1 call for sealed + 1 call for growing = 2
		s.Equal(int32(3), callCount.Load())
	})
}

func (s *TwoStageSearchSuite) TestTwoStageSearchZeroValidCounts() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("zero_valid_counts_all_segments", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		// Return zero valid counts - no records pass the filter
		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Return(&internalpb.SearchResults{
				FilterValidCounts: []int64{0, 0, 0},
			}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments
		s.delegator.distribution.AddDistributions(
			SegmentEntry{NodeID: 1, SegmentID: 7000},
			SegmentEntry{NodeID: 1, SegmentID: 7001},
			SegmentEntry{NodeID: 1, SegmentID: 7002},
		)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 7000, NodeID: 1},
					{SegmentID: 7001, NodeID: 1},
					{SegmentID: 7002, NodeID: 1},
				},
			},
		}
		growing := []SegmentEntry{}
		sealedRowCount := map[int64]int64{
			7000: 10000,
			7001: 10000,
			7002: 10000,
		}

		results, fallback, err := s.delegator.twoStageSearch(ctx, req, sealed, growing, sealedRowCount)
		s.NoError(err)
		s.False(fallback)
		s.NotNil(results)
	})
}

func (s *TwoStageSearchSuite) TestExecuteFilterStageWithMultipleNodes() {
	s.delegator.Start()
	paramtable.SetNodeID(1)

	s.Run("multiple_nodes_success", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}
		worker3 := &cluster.MockWorker{}
		workers[1] = worker1
		workers[2] = worker2
		workers[3] = worker3

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.True(req.GetFilterOnly(), "FilterOnly should be true")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{100},
		}, nil)

		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.True(req.GetFilterOnly(), "FilterOnly should be true")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{200},
		}, nil)

		worker3.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.True(req.GetFilterOnly(), "FilterOnly should be true")
			}).Return(&internalpb.SearchResults{
			FilterValidCounts: []int64{300},
		}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		// Setup distribution with segments on multiple nodes
		s.delegator.distribution.AddDistributions(
			SegmentEntry{NodeID: 1, SegmentID: 8000},
			SegmentEntry{NodeID: 2, SegmentID: 8001},
			SegmentEntry{NodeID: 3, SegmentID: 8002},
		)

		ctx := context.Background()
		req := &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Topk: 100,
			},
			DmlChannels: []string{s.vchannelName},
		}

		sealed := []SnapshotItem{
			{
				NodeID: 1,
				Segments: []SegmentEntry{
					{SegmentID: 8000, NodeID: 1},
				},
			},
			{
				NodeID: 2,
				Segments: []SegmentEntry{
					{SegmentID: 8001, NodeID: 2},
				},
			},
			{
				NodeID: 3,
				Segments: []SegmentEntry{
					{SegmentID: 8002, NodeID: 3},
				},
			},
		}
		sealedRowCount := map[int64]int64{
			8000: 10000,
			8001: 10000,
			8002: 10000,
		}

		validCounts, err := s.delegator.executeFilterStage(ctx, req, sealed, sealedRowCount)
		s.NoError(err)
		s.NotNil(validCounts)
		s.Len(validCounts, 3)

		// Verify total valid counts
		totalValidCount := int64(0)
		for _, count := range validCounts {
			totalValidCount += count
		}
		s.Equal(int64(600), totalValidCount)
	})
}

func TestTwoStageSearchSuite(t *testing.T) {
	suite.Run(t, new(TwoStageSearchSuite))
}
