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

import (
	"context"
	"testing"
	"time"

	"github.com/bytedance/mockey"
	"github.com/cockroachdb/errors"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
	"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
	"github.com/milvus-io/milvus/internal/querycoordv2/params"
	"github.com/milvus-io/milvus/internal/querycoordv2/session"
	"github.com/milvus-io/milvus/internal/querycoordv2/utils"
	"github.com/milvus-io/milvus/pkg/v2/kv"
	"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
	"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v2/util/etcd"
	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)

type ExecutorTestSuite struct {
	suite.Suite

	// Dependencies
	kv      kv.MetaKv
	meta    *meta.Meta
	dist    *meta.DistributionManager
	target  *meta.TargetManager
	broker  *meta.MockBroker
	nodeMgr *session.NodeManager
	cluster *session.MockCluster

	// Test object
	executor *Executor
	ctx      context.Context
}

func (suite *ExecutorTestSuite) SetupSuite() {
	paramtable.Init()
}

func (suite *ExecutorTestSuite) SetupTest() {
	var err error
	config := params.GenerateEtcdConfig()
	cli, err := etcd.GetEtcdClient(
		config.UseEmbedEtcd.GetAsBool(),
		config.EtcdUseSSL.GetAsBool(),
		config.Endpoints.GetAsStrings(),
		config.EtcdTLSCert.GetValue(),
		config.EtcdTLSKey.GetValue(),
		config.EtcdTLSCACert.GetValue(),
		config.EtcdTLSMinVersion.GetValue())
	suite.Require().NoError(err)
	suite.kv = etcdkv.NewEtcdKV(cli, config.MetaRootPath.GetValue())

	// Initialize components
	store := querycoord.NewCatalog(suite.kv)
	idAllocator := params.RandomIncrementIDAllocator()
	suite.nodeMgr = session.NewNodeManager()
	suite.meta = meta.NewMeta(idAllocator, store, suite.nodeMgr)
	suite.dist = meta.NewDistributionManager()
	suite.broker = meta.NewMockBroker(suite.T())
	suite.target = meta.NewTargetManager(suite.broker, suite.meta)
	suite.cluster = session.NewMockCluster(suite.T())

	suite.executor = NewExecutor(suite.meta, suite.dist, suite.broker, suite.target, suite.cluster, suite.nodeMgr)
	suite.ctx = context.Background()
}

func (suite *ExecutorTestSuite) TearDownTest() {
	suite.kv.Close()
}

func (suite *ExecutorTestSuite) TestReleaseSegmentWithServiceableLeader() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create segment task
	action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
	suite.NoError(err)

	// Setup serviceable leader view (UnServiceableError = nil means serviceable)
	serviceableLeaderView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
	serviceableLeaderView.UnServiceableError = nil // serviceable
	suite.dist.LeaderViewManager.Update(1, serviceableLeaderView)

	// Setup non-serviceable leader view
	nonServiceableLeaderView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
	nonServiceableLeaderView.UnServiceableError = errors.New("not serviceable") // not serviceable
	suite.dist.LeaderViewManager.Update(2, nonServiceableLeaderView)

	// Mock cluster response
	suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(1), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute release segment
	suite.executor.releaseSegment(task, 0)

	// Verify that serviceable leader (node 1) was chosen for release
	suite.cluster.AssertExpectations(suite.T())
}

func (suite *ExecutorTestSuite) TestReleaseSegmentFallbackToNonServiceableLeader() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create segment task
	action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
	suite.NoError(err)

	// Setup only non-serviceable leader view
	nonServiceableLeaderView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
	nonServiceableLeaderView.UnServiceableError = errors.New("not serviceable") // not serviceable
	suite.dist.LeaderViewManager.Update(2, nonServiceableLeaderView)

	// Mock cluster response for fallback to non-serviceable leader
	suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute release segment
	suite.executor.releaseSegment(task, 0)

	// Verify that non-serviceable leader (node 2) was used as fallback
	suite.cluster.AssertExpectations(suite.T())
}

func (suite *ExecutorTestSuite) TestReleaseSegmentNoLeaderFound() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create segment task
	action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
	suite.NoError(err)

	// No leader views setup - should trigger early return

	// Execute release segment - should return early without calling cluster
	suite.executor.releaseSegment(task, 0)

	// Verify that no cluster calls were made
	suite.cluster.AssertExpectations(suite.T())
}

func (suite *ExecutorTestSuite) TestReleaseSegmentWithNodeNotInReplica() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   3, // Node 3 not in replica
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 3)

	// Create segment task with node not in replica
	action := NewSegmentAction(3, ActionTypeReduce, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
	suite.NoError(err)

	// Mock cluster response for direct node release
	suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(3), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute release segment
	suite.executor.releaseSegment(task, 0)

	// Verify that direct node release was used
	suite.cluster.AssertExpectations(suite.T())
}

func (suite *ExecutorTestSuite) TestReleaseSegmentChannelSpecificLookup() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create segment task
	action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
	suite.NoError(err)

	// Setup leader view for a different channel (should not be found with channel filter)
	leaderViewDifferentChannel := utils.CreateTestLeaderView(1, 1, "other-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
	leaderViewDifferentChannel.UnServiceableError = errors.New("not serviceable")
	suite.dist.LeaderViewManager.Update(1, leaderViewDifferentChannel)

	// Setup leader view for the correct channel but non-serviceable
	leaderViewCorrectChannel := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
	leaderViewCorrectChannel.UnServiceableError = errors.New("not serviceable")
	suite.dist.LeaderViewManager.Update(2, leaderViewCorrectChannel)

	// Mock cluster response
	suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute release segment
	suite.executor.releaseSegment(task, 0)

	// Verify that channel-specific lookup worked correctly
	suite.cluster.AssertExpectations(suite.T())
}

func (suite *ExecutorTestSuite) TestBalanceTaskWithTwoDelegators() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create balance task with load and release actions
	loadAction := NewSegmentAction(2, ActionTypeGrow, "test-channel", 100)      // Load on node 2
	releaseAction := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100) // Release from node 1

	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, loadAction, releaseAction)
	suite.NoError(err)

	// Setup old delegator (node 1) - serviceable
	oldDelegatorView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
	oldDelegatorView.Version = 1
	oldDelegatorView.UnServiceableError = nil // serviceable
	suite.dist.LeaderViewManager.Update(1, oldDelegatorView)

	// Mock broker responses
	suite.broker.EXPECT().DescribeCollection(mock.Anything, int64(1)).Return(&milvuspb.DescribeCollectionResponse{
		Schema: &schemapb.CollectionSchema{
			Name: "TestBalanceTask",
			Fields: []*schemapb.FieldSchema{
				{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
			},
		},
	}, nil)
	suite.broker.EXPECT().GetSegmentInfo(mock.Anything, int64(100)).Return([]*datapb.SegmentInfo{
		{
			ID:            100,
			CollectionID:  1,
			PartitionID:   1,
			InsertChannel: "test-channel",
		},
	}, nil)
	suite.broker.EXPECT().GetIndexInfo(mock.Anything, int64(1), int64(100)).Return(nil, nil)
	suite.broker.EXPECT().ListIndexes(mock.Anything, int64(1)).Return([]*indexpb.IndexInfo{
		{
			CollectionID: 1,
		},
	}, nil)

	// Setup target for collection to be loaded
	channel := &datapb.VchannelInfo{
		CollectionID: 1,
		ChannelName:  "test-channel",
	}
	segments := []*datapb.SegmentInfo{
		{
			ID:            100,
			CollectionID:  1,
			PartitionID:   1,
			InsertChannel: "test-channel",
		},
	}
	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return([]*datapb.VchannelInfo{channel}, segments, nil)
	suite.target.UpdateCollectionNextTarget(suite.ctx, 1)

	// Expect load to be called on latest serviceable delegator (node 1)
	suite.cluster.EXPECT().LoadSegments(mock.Anything, int64(1), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute load segment (step 0)
	suite.executor.loadSegment(task, 0)

	// Setup new delegator (node 2) - serviceable
	newDelegatorView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
	newDelegatorView.Version = 2
	newDelegatorView.UnServiceableError = nil // serviceable
	suite.dist.LeaderViewManager.Update(2, newDelegatorView)

	// Execute release segment (step 1)
	suite.executor.releaseSegment(task, 1)

	// verify that the task is failed due to shard leader change
	suite.Error(task.Err())
}

func (suite *ExecutorTestSuite) TestBalanceTaskFallbackToLatestWhenNoServiceableDelegator() {
	// Setup collection and replica
	collection := utils.CreateTestCollection(1, 1)
	suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
	suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))

	replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
	suite.meta.ReplicaManager.Put(suite.ctx, replica)

	// Setup nodes
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
	suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

	// Create load action
	loadAction := NewSegmentAction(2, ActionTypeGrow, "test-channel", 100)
	task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, loadAction)
	suite.NoError(err)

	// Setup old delegator (node 1) - not serviceable
	oldDelegatorView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
	oldDelegatorView.UnServiceableError = errors.New("not serviceable")
	oldDelegatorView.Version = 1
	suite.dist.LeaderViewManager.Update(1, oldDelegatorView)

	// Setup new delegator (node 2) - not serviceable but latest
	newDelegatorView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
	newDelegatorView.UnServiceableError = errors.New("not serviceable")
	newDelegatorView.Version = 2 // latest version
	suite.dist.LeaderViewManager.Update(2, newDelegatorView)

	// Mock broker responses
	suite.broker.EXPECT().DescribeCollection(mock.Anything, int64(1)).Return(&milvuspb.DescribeCollectionResponse{
		Schema: &schemapb.CollectionSchema{
			Name: "TestBalanceTask",
			Fields: []*schemapb.FieldSchema{
				{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
			},
		},
	}, nil)
	suite.broker.EXPECT().GetSegmentInfo(mock.Anything, int64(100)).Return([]*datapb.SegmentInfo{
		{
			ID:            100,
			CollectionID:  1,
			PartitionID:   1,
			InsertChannel: "test-channel",
		},
	}, nil)
	suite.broker.EXPECT().GetIndexInfo(mock.Anything, int64(1), int64(100)).Return(nil, nil)
	suite.broker.EXPECT().ListIndexes(mock.Anything, int64(1)).Return([]*indexpb.IndexInfo{
		{
			CollectionID: 1,
		},
	}, nil)

	// Setup target for collection to be loaded
	channel := &datapb.VchannelInfo{
		CollectionID: 1,
		ChannelName:  "test-channel",
	}
	segments := []*datapb.SegmentInfo{
		{
			ID:            100,
			CollectionID:  1,
			PartitionID:   1,
			InsertChannel: "test-channel",
		},
	}
	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return([]*datapb.VchannelInfo{channel}, segments, nil)
	suite.target.UpdateCollectionNextTarget(suite.ctx, 1)

	// Expect load to be called on the latest delegator (node 2) as fallback
	suite.cluster.EXPECT().LoadSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
	}, nil).Once()

	// Execute load segment
	suite.executor.loadSegment(task, 0)

	// Verify that the latest delegator was chosen as fallback
	suite.cluster.AssertExpectations(suite.T())
}

func TestExecutorSuite(t *testing.T) {
	suite.Run(t, new(ExecutorTestSuite))
}

// L0 segment specific tests for delegator selection during load
func (suite *ExecutorTestSuite) TestLoadL0_SelectDelegatorMissingSegment() {
	t := suite.T()
	mockey.PatchConvey("TestLoadL0_SelectDelegatorMissingSegment", t, func() {
		// Setup collection and replica
		collectionID := int64(1)
		segmentID := int64(100)
		channelName := "test-channel"

		collection := utils.CreateTestCollection(collectionID, 1)
		suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
		suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))

		replica := utils.CreateTestReplica(1, collectionID, []int64{1, 2})
		suite.meta.ReplicaManager.Put(suite.ctx, replica)

		// Setup nodes
		suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
		suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2, Address: "localhost", Hostname: "localhost"}))
		suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
		suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

		// Only one delegator view which is missing the L0 segment
		viewMissing := utils.CreateTestLeaderView(2, collectionID, channelName, map[int64]int64{}, map[int64]*meta.Segment{})
		viewMissing.UnServiceableError = nil
		viewMissing.Version = 1
		suite.dist.LeaderViewManager.Update(2, viewMissing)

		// Mock broker responses using mockey
		mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
			ID:            segmentID,
			CollectionID:  collectionID,
			PartitionID:   1,
			InsertChannel: channelName,
			Level:         datapb.SegmentLevel_L0,
		}}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
		mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()

		channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
		segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
		mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
		suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)

		// Capture LoadSegments target node
		var capturedNode int64 = -1
		mockey.Mock((*session.MockCluster).LoadSegments).To(func(m *session.MockCluster, ctx context.Context, node int64, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
			capturedNode = node
			return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
		}).Build()

		// Execute load segment
		action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
		task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
		suite.NoError(err)
		_ = suite.executor.loadSegment(task, 0)

		suite.Equal(int64(2), capturedNode)
	})
}

func (suite *ExecutorTestSuite) TestLoadL0_FallbackWhenAllDelegatorsHaveSegment() {
	t := suite.T()
	mockey.PatchConvey("TestLoadL0_FallbackWhenAllDelegatorsHaveSegment", t, func() {
		// Setup collection and replica
		collectionID := int64(1)
		segmentID := int64(101)
		channelName := "test-channel"

		collection := utils.CreateTestCollection(collectionID, 1)
		suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
		suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))

		replica := utils.CreateTestReplica(1, collectionID, []int64{1, 2})
		suite.meta.ReplicaManager.Put(suite.ctx, replica)

		// Setup nodes
		suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
		suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2, Address: "localhost", Hostname: "localhost"}))
		suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
		suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)

		// Both delegators already have the L0 segment; should fallback to serviceable latest
		view1 := utils.CreateTestLeaderView(1, collectionID, channelName, map[int64]int64{segmentID: 1}, map[int64]*meta.Segment{})
		view1.UnServiceableError = nil
		view1.Version = 2
		suite.dist.LeaderViewManager.Update(1, view1)

		view2 := utils.CreateTestLeaderView(2, collectionID, channelName, map[int64]int64{segmentID: 2}, map[int64]*meta.Segment{})
		view2.UnServiceableError = errors.New("not serviceable")
		view2.Version = 3
		suite.dist.LeaderViewManager.Update(2, view2)

		// Mock broker responses using mockey
		mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
			ID:            segmentID,
			CollectionID:  collectionID,
			PartitionID:   1,
			InsertChannel: channelName,
			Level:         datapb.SegmentLevel_L0,
		}}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
		mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()

		channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
		segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
		mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
		suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)

		// Capture LoadSegments target node
		var capturedNode int64 = -1
		mockey.Mock((*session.MockCluster).LoadSegments).To(func(m *session.MockCluster, ctx context.Context, node int64, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
			capturedNode = node
			return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
		}).Build()

		action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
		task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
		suite.NoError(err)
		_ = suite.executor.loadSegment(task, 0)

		suite.Equal(int64(1), capturedNode)
	})
}

func (suite *ExecutorTestSuite) TestLoadL0_NoDelegatorFound() {
	t := suite.T()
	mockey.PatchConvey("TestLoadL0_NoDelegatorFound", t, func() {
		// Setup collection and replica
		collectionID := int64(1)
		segmentID := int64(102)
		channelName := "test-channel"

		collection := utils.CreateTestCollection(collectionID, 1)
		suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
		suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))

		replica := utils.CreateTestReplica(1, collectionID, []int64{1})
		suite.meta.ReplicaManager.Put(suite.ctx, replica)

		// Setup node
		suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
		suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)

		// No leader views set for the channel -> should return error early

		// Mock broker responses using mockey
		mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
			ID:            segmentID,
			CollectionID:  collectionID,
			PartitionID:   1,
			InsertChannel: channelName,
			Level:         datapb.SegmentLevel_L0,
		}}, nil).Build()
		mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
		mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()

		channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
		segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
		mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
		suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)

		action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
		task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
		suite.NoError(err)
		err = suite.executor.loadSegment(task, 0)
		suite.Error(err)
	})
}