package replicates

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/metricsutil"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/recovery"
	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
	"github.com/milvus-io/milvus/pkg/v3/streaming/walimpls/impls/walimplstest"
	"github.com/milvus-io/milvus/pkg/v3/util/replicateutil"
)

func TestNonReplicateManager(t *testing.T) {
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID:           walimplstest.NewTestMessageID(1),
				TimeTick:            1,
				ReplicateCheckpoint: nil,
				ReplicateConfig:     nil,
			},
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)

	testSwitchReplicateMode(t, rm, "test1", "test2")
	testMessageOnPrimary(t, rm)
	testMessageOnSecondary(t, rm)
}

func TestPrimaryReplicateManager(t *testing.T) {
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID:           walimplstest.NewTestMessageID(1),
				TimeTick:            1,
				ReplicateCheckpoint: nil,
				ReplicateConfig:     newReplicateConfiguration("test1", "test2"),
			},
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)

	testSwitchReplicateMode(t, rm, "test1", "test2")
	testMessageOnPrimary(t, rm)
	testMessageOnSecondary(t, rm)
}

func TestSecondaryReplicateManager(t *testing.T) {
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(1),
					TimeTick:  1,
				},
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
			},
			TxnBuffer: utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics()),
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)

	testSwitchReplicateMode(t, rm, "test1", "test2")
	testMessageOnPrimary(t, rm)
	testMessageOnSecondary(t, rm)
}

func TestReplicateManager_UnreplicableMessage(t *testing.T) {
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID:       walimplstest.NewTestMessageID(1),
				TimeTick:        1,
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(1),
					TimeTick:  1,
				},
			},
			TxnBuffer: utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics()),
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())

	msg := newCreateSnapshotMutableMessage()
	g, err := rm.BeginReplicateMessage(context.Background(), msg)
	assert.ErrorIs(t, err, ErrNotHandledByReplicateManager)
	assert.Nil(t, g)

	msg = newCreateSnapshotMutableMessage().WithReplicateHeader(&message.ReplicateHeader{
		ClusterID:              "test2",
		MessageID:              walimplstest.NewTestMessageID(2),
		LastConfirmedMessageID: walimplstest.NewTestMessageID(1),
		TimeTick:               2,
		VChannel:               "test2-rootcoord-dml_0",
	})
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
	assert.Nil(t, g)
}

func TestSalvageCheckpointCaptureOnForcePromote(t *testing.T) {
	// Setup: cluster starts as secondary with a checkpoint
	txnBuffer := utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(100),
					TimeTick:  1000,
				},
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
			},
			TxnBuffer: txnBuffer,
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())

	// Initially no salvage checkpoint
	assert.Empty(t, rm.GetSalvageCheckpoint())

	// Force promote to primary
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "test2"))
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())

	// Salvage checkpoint should be captured
	salvageCPs := rm.GetSalvageCheckpoint()
	assert.Len(t, salvageCPs, 1)
	salvageCP := salvageCPs[0]
	assert.Equal(t, "test2", salvageCP.ClusterID)
	assert.Equal(t, "test2-rootcoord-dml_0", salvageCP.PChannel)
	assert.Equal(t, uint64(1000), salvageCP.TimeTick)
}

func TestSalvageCheckpointNotCapturedOnNormalPromote(t *testing.T) {
	// Setup: cluster starts as secondary
	txnBuffer := utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(100),
					TimeTick:  1000,
				},
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
			},
			TxnBuffer: txnBuffer,
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())

	// Normal promote (not force)
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", false, "test2"))
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())

	// Salvage checkpoint should NOT be captured
	assert.Empty(t, rm.GetSalvageCheckpoint())
}

func TestSalvageCheckpointNotCapturedWhenAlreadyPrimary(t *testing.T) {
	// Setup: cluster starts as primary (no secondary state)
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID:           walimplstest.NewTestMessageID(1),
				TimeTick:            1,
				ReplicateCheckpoint: nil,
				ReplicateConfig:     newReplicateConfiguration("test1", "test2"),
			},
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())

	// Force promote when already primary (no secondary state)
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "test2"))
	assert.NoError(t, err)

	// Salvage checkpoint should NOT be captured (no secondary state to capture)
	assert.Empty(t, rm.GetSalvageCheckpoint())
}

func TestSalvageCheckpointLoadedFromEtcd(t *testing.T) {
	// Simulate WAL recovery where salvage checkpoints were previously persisted to etcd
	// and loaded back into the ReplicateManagerRecoverParam.
	preLoaded := []*utility.ReplicateCheckpoint{
		{ClusterID: "cluster-a", PChannel: "cluster-a-rootcoord-dml_0", TimeTick: 100},
		{ClusterID: "cluster-b", PChannel: "cluster-b-rootcoord-dml_0", TimeTick: 200},
	}

	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID:       walimplstest.NewTestMessageID(1),
				TimeTick:        1,
				ReplicateConfig: newReplicateConfiguration("test1", "test2"),
			},
		},
		SalvageCheckpoints: preLoaded,
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())

	salvageCPs := rm.GetSalvageCheckpoint()
	assert.Len(t, salvageCPs, 2)

	byCluster := make(map[string]*utility.ReplicateCheckpoint, 2)
	for _, cp := range salvageCPs {
		byCluster[cp.ClusterID] = cp
	}
	assert.Equal(t, uint64(100), byCluster["cluster-a"].TimeTick)
	assert.Equal(t, "cluster-a-rootcoord-dml_0", byCluster["cluster-a"].PChannel)
	assert.Equal(t, uint64(200), byCluster["cluster-b"].TimeTick)
	assert.Equal(t, "cluster-b-rootcoord-dml_0", byCluster["cluster-b"].PChannel)
}

func TestSalvageCheckpointMultipleForcePromotes(t *testing.T) {
	// Start as secondary of cluster-a, force promote to primary,
	// then become secondary of cluster-b, force promote again.
	// Both salvage checkpoints should accumulate (keyed by source cluster).
	txnBuffer := utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "cluster-a",
					PChannel:  "cluster-a-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(10),
					TimeTick:  100,
				},
				ReplicateConfig: newReplicateConfiguration("cluster-a", "test1"),
			},
			TxnBuffer: txnBuffer,
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())
	assert.Empty(t, rm.GetSalvageCheckpoint())

	// First force promote: secondary(cluster-a) -> primary
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "cluster-a"))
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())
	cps := rm.GetSalvageCheckpoint()
	assert.Len(t, cps, 1)
	assert.Equal(t, "cluster-a", cps[0].ClusterID)

	// Now become secondary of cluster-b
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("cluster-b", "test1"))
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())

	// Still one salvage checkpoint from first promote
	assert.Len(t, rm.GetSalvageCheckpoint(), 1)

	// Second force promote: secondary(cluster-b) -> primary
	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "cluster-b"))
	assert.NoError(t, err)
	assert.Equal(t, replicateutil.RolePrimary, rm.Role())

	// Both salvage checkpoints should now be present
	cps = rm.GetSalvageCheckpoint()
	assert.Len(t, cps, 2)
	byCluster := make(map[string]*utility.ReplicateCheckpoint, 2)
	for _, cp := range cps {
		byCluster[cp.ClusterID] = cp
	}
	assert.Contains(t, byCluster, "cluster-a")
	assert.Contains(t, byCluster, "cluster-b")
}

func TestSecondaryReplicateManagerWithTxn(t *testing.T) {
	txnBuffer := utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
	txnMsgs := newReplicateTxnMessage("test1", "test2", 2)

	for _, msg := range txnMsgs[0:3] {
		immutableMsg := msg.WithTimeTick(3).IntoImmutableMessage(walimplstest.NewTestMessageID(1))
		txnBuffer.HandleImmutableMessages([]message.ImmutableMessage{immutableMsg}, msg.TimeTick())
	}

	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(1),
					TimeTick:  1,
				},
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
			},
			TxnBuffer: txnBuffer,
		},
	})
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)

	committed := false
	for _, msg := range newReplicateTxnMessage("test1", "test2", 2) {
		g, err := rm.BeginReplicateMessage(context.Background(), msg)
		if msg.MessageType() == message.MessageTypeCommitTxn && !committed {
			assert.NoError(t, err)
			assert.NotNil(t, g)
			g.Ack(nil)
			committed = true
		} else {
			assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
			assert.Nil(t, g)
		}
	}
}

func TestSecondaryReplicateManagerIgnoreStaleTxnBody(t *testing.T) {
	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
		ChannelInfo: types.PChannelInfo{
			Name: "test1-rootcoord-dml_0",
			Term: 1,
		},
		CurrentClusterID: "test1",
		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
			Checkpoint: &utility.WALCheckpoint{
				MessageID: walimplstest.NewTestMessageID(1),
				TimeTick:  1,
				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
					ClusterID: "test2",
					PChannel:  "test2-rootcoord-dml_0",
					MessageID: walimplstest.NewTestMessageID(1),
					TimeTick:  100,
				},
				ReplicateConfig: newReplicateConfiguration("test2", "test1"),
			},
			TxnBuffer: utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics()),
		},
	})
	assert.NoError(t, err)

	beginCurrentTxn := newReplicateTxnMessageWithTxnID("test1", "test2", 200, 2)[0]
	g, err := rm.BeginReplicateMessage(context.Background(), beginCurrentTxn)
	assert.NoError(t, err)
	assert.NotNil(t, g)
	g.Ack(nil)

	staleTxnBody := newReplicateTxnMessageWithTxnID("test1", "test2", 100, 1)[2]
	g, err = rm.BeginReplicateMessage(context.Background(), staleTxnBody)
	assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
	assert.Nil(t, g)
}

func testSwitchReplicateMode(t *testing.T, rm ReplicatesManager, primaryClusterID, secondaryClusterID string) {
	ctx := context.Background()

	// switch to primary
	err := rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err := rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// idempotent switch to primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// switch to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// idempotent switch to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// switch back to primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// idempotent switch back to primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// switch back to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// idempotent switch back to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// add a new cluster and switch to primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID, "test3"))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// idempotent add a new cluster and switch to primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(primaryClusterID, secondaryClusterID, "test3"))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RolePrimary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, cp)

	// add a new cluster and switch to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID, "test3"))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// idempotent add a new cluster and switch to secondary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage(secondaryClusterID, primaryClusterID, "test3"))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, secondaryClusterID)
	assert.Equal(t, cp.PChannel, secondaryClusterID+"-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// switch the primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage("test3", primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, "test3")
	assert.Equal(t, cp.PChannel, "test3-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))

	// idempotent switch the primary
	err = rm.SwitchReplicateMode(ctx, newAlterReplicateConfigMessage("test3", primaryClusterID, secondaryClusterID))
	assert.NoError(t, err)
	assert.Equal(t, rm.Role(), replicateutil.RoleSecondary)
	cp, err = rm.GetReplicateCheckpoint()
	assert.NoError(t, err)
	assert.Equal(t, cp.ClusterID, "test3")
	assert.Equal(t, cp.PChannel, "test3-rootcoord-dml_0")
	assert.Nil(t, cp.MessageID)
	assert.Equal(t, cp.TimeTick, uint64(0))
}

func testMessageOnPrimary(t *testing.T, rm ReplicatesManager) {
	// switch to primary
	err := rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("test1", "test2"))
	assert.NoError(t, err)

	// Test self-controlled message
	g, err := rm.BeginReplicateMessage(context.Background(), message.NewCreateSegmentMessageBuilderV2().
		WithHeader(&message.CreateSegmentMessageHeader{}).
		WithBody(&message.CreateSegmentMessageBody{}).
		WithVChannel("test1-rootcoord-dml_0").
		MustBuildMutable())
	assert.ErrorIs(t, err, ErrNotHandledByReplicateManager)
	assert.Nil(t, g)

	// Test non-replicate message
	msg := newNonReplicateMessage("test1")
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.ErrorIs(t, err, ErrNotHandledByReplicateManager)
	assert.Nil(t, g)

	// Test replicate message
	msg = newReplicateMessage("test1", "test2")
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, g)
}

func testMessageOnSecondary(t *testing.T, rm ReplicatesManager) {
	// switch to secondary
	err := rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("test2", "test1"))
	assert.NoError(t, err)

	// Test wrong cluster replicates
	msg := newReplicateMessage("test1", "test3")
	g, err := rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, g)

	// Test self-controlled message
	g, err = rm.BeginReplicateMessage(context.Background(), message.NewCreateSegmentMessageBuilderV2().
		WithHeader(&message.CreateSegmentMessageHeader{}).
		WithBody(&message.CreateSegmentMessageBody{}).
		WithVChannel("test1-rootcoord-dml_0").
		MustBuildMutable())
	assert.ErrorIs(t, err, ErrNotHandledByReplicateManager)
	assert.Nil(t, g)

	// Test non-replicate message
	msg = newNonReplicateMessage("test1")
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsReplicateViolation())
	assert.Nil(t, g)

	// Test replicate message
	msg = newReplicateMessage("test1", "test2")
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.NoError(t, err)
	assert.NotNil(t, g)
	g.Ack(nil)

	// Test replicate message
	msg = newReplicateMessage("test1", "test2")
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
	assert.Nil(t, g)

	for idx, msg := range newReplicateTxnMessage("test1", "test2", 2) {
		g, err = rm.BeginReplicateMessage(context.Background(), msg)
		if idx%2 == 0 {
			assert.NoError(t, err)
			assert.NotNil(t, g)
			g.Ack(nil)
		} else {
			assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
			assert.Nil(t, g)
		}
	}

	msg = newReplicateMessage("test1", "test2", 2)
	g, err = rm.BeginReplicateMessage(context.Background(), msg)
	assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
	assert.Nil(t, g)

	g, err = rm.BeginReplicateMessage(context.Background(), newReplicateTxnMessage("test1", "test2", 2)[0])
	assert.True(t, status.AsStreamingError(err).IsIgnoredOperation())
	assert.Nil(t, g)
}

func newCreateSnapshotMutableMessage() message.MutableMessage {
	return message.NewCreateSnapshotMessageBuilderV2().
		WithHeader(&message.CreateSnapshotMessageHeader{}).
		WithBody(&message.CreateSnapshotMessageBody{}).
		WithVChannel("test1-rootcoord-dml_0").
		WithUnreplicable().
		MustBuildMutable()
}

// newReplicateConfiguration creates a valid replicate configuration for testing
func newReplicateConfiguration(primaryClusterID string, secondaryClusterID ...string) *commonpb.ReplicateConfiguration {
	clusters := []*commonpb.MilvusCluster{
		{ClusterId: primaryClusterID, Pchannels: []string{primaryClusterID + "-rootcoord-dml_0", primaryClusterID + "-rootcoord-dml_1"}},
	}
	crossClusterTopology := []*commonpb.CrossClusterTopology{}
	for _, secondaryClusterID := range secondaryClusterID {
		clusters = append(clusters, &commonpb.MilvusCluster{ClusterId: secondaryClusterID, Pchannels: []string{secondaryClusterID + "-rootcoord-dml_0", secondaryClusterID + "-rootcoord-dml_1"}})
		crossClusterTopology = append(crossClusterTopology, &commonpb.CrossClusterTopology{SourceClusterId: primaryClusterID, TargetClusterId: secondaryClusterID})
	}
	return &commonpb.ReplicateConfiguration{
		Clusters:             clusters,
		CrossClusterTopology: crossClusterTopology,
	}
}

func newAlterReplicateConfigMessage(primaryClusterID string, secondaryClusterID ...string) message.MutableAlterReplicateConfigMessageV2 {
	return newAlterReplicateConfigMessageWithForcePromote(primaryClusterID, false, secondaryClusterID...)
}

func newAlterReplicateConfigMessageWithForcePromote(primaryClusterID string, forcePromote bool, secondaryClusterID ...string) message.MutableAlterReplicateConfigMessageV2 {
	return message.MustAsMutableAlterReplicateConfigMessageV2(message.NewAlterReplicateConfigMessageBuilderV2().
		WithHeader(&message.AlterReplicateConfigMessageHeader{
			ReplicateConfiguration: newReplicateConfiguration(primaryClusterID, secondaryClusterID...),
			ForcePromote:           forcePromote,
		}).
		WithBody(&message.AlterReplicateConfigMessageBody{}).
		WithVChannel(primaryClusterID + "-rootcoord-dml_0").
		MustBuildMutable())
}

func newNonReplicateMessage(clusterID string) message.MutableMessage {
	return message.NewCreateDatabaseMessageBuilderV2().
		WithHeader(&message.CreateDatabaseMessageHeader{}).
		WithBody(&message.CreateDatabaseMessageBody{}).
		WithVChannel(clusterID + "-rootcoord-dml_0").
		MustBuildMutable()
}

func newReplicateMessage(clusterID string, sourceClusterID string, timetick ...uint64) message.MutableMessage {
	tt := uint64(1)
	if len(timetick) > 0 {
		tt = timetick[0]
	}
	msg := message.NewCreateDatabaseMessageBuilderV2().
		WithHeader(&message.CreateDatabaseMessageHeader{}).
		WithBody(&message.CreateDatabaseMessageBody{}).
		WithVChannel(sourceClusterID + "-rootcoord-dml_0").
		MustBuildMutable().
		WithTimeTick(tt).
		WithLastConfirmed(walimplstest.NewTestMessageID(1)).
		IntoImmutableMessage(walimplstest.NewTestMessageID(1))

	replicateMsg := message.MustNewReplicateMessage(
		sourceClusterID,
		msg.IntoImmutableMessageProto(),
	)
	replicateMsg.OverwriteReplicateVChannel(
		clusterID + "-rootcoord-dml_0",
	)
	return replicateMsg
}

func newImmutableTxnMessage(clusterID string, timetick ...uint64) []message.ImmutableMessage {
	tt := uint64(1)
	if len(timetick) > 0 {
		tt = timetick[0]
	}
	immutables := []message.ImmutableMessage{
		message.NewBeginTxnMessageBuilderV2().
			WithHeader(&message.BeginTxnMessageHeader{}).
			WithBody(&message.BeginTxnMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(message.TxnContext{
				TxnID:     message.TxnID(1),
				Keepalive: message.TxnKeepaliveInfinite,
			}).
			WithTimeTick(tt).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
		message.NewCreateDatabaseMessageBuilderV2().
			WithHeader(&message.CreateDatabaseMessageHeader{}).
			WithBody(&message.CreateDatabaseMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(message.TxnContext{
				TxnID:     message.TxnID(1),
				Keepalive: message.TxnKeepaliveInfinite,
			}).
			WithTimeTick(tt).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
		message.NewCommitTxnMessageBuilderV2().
			WithHeader(&message.CommitTxnMessageHeader{}).
			WithBody(&message.CommitTxnMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(message.TxnContext{
				TxnID:     message.TxnID(1),
				Keepalive: message.TxnKeepaliveInfinite,
			}).
			WithTimeTick(tt).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
	}
	return immutables
}

func newReplicateTxnMessage(clusterID string, sourceClusterID string, timetick ...uint64) []message.MutableMessage {
	tt := uint64(1)
	if len(timetick) > 0 {
		tt = timetick[0]
	}
	return newReplicateTxnMessageWithTxnID(clusterID, sourceClusterID, tt, 1)
}

func newReplicateTxnMessageWithTxnID(clusterID string, sourceClusterID string, timetick uint64, txnID message.TxnID) []message.MutableMessage {
	immutables := newImmutableTxnMessageWithTxnID(sourceClusterID, timetick, txnID)
	replicateMsgs := []message.MutableMessage{}
	for _, immutable := range immutables {
		replicateMsg := message.MustNewReplicateMessage(
			sourceClusterID,
			immutable.IntoImmutableMessageProto(),
		)
		replicateMsg.OverwriteReplicateVChannel(
			clusterID + "-rootcoord-dml_0",
		)
		replicateMsgs = append(replicateMsgs, replicateMsg)
		// test the idempotency
		replicateMsgs = append(replicateMsgs, replicateMsg)
	}
	return replicateMsgs
}

func newImmutableTxnMessageWithTxnID(clusterID string, timetick uint64, txnID message.TxnID) []message.ImmutableMessage {
	txnCtx := message.TxnContext{
		TxnID:     txnID,
		Keepalive: message.TxnKeepaliveInfinite,
	}
	immutables := []message.ImmutableMessage{
		message.NewBeginTxnMessageBuilderV2().
			WithHeader(&message.BeginTxnMessageHeader{}).
			WithBody(&message.BeginTxnMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(txnCtx).
			WithTimeTick(timetick).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
		message.NewCreateDatabaseMessageBuilderV2().
			WithHeader(&message.CreateDatabaseMessageHeader{}).
			WithBody(&message.CreateDatabaseMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(txnCtx).
			WithTimeTick(timetick).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
		message.NewCommitTxnMessageBuilderV2().
			WithHeader(&message.CommitTxnMessageHeader{}).
			WithBody(&message.CommitTxnMessageBody{}).
			WithVChannel(clusterID + "-rootcoord-dml_0").
			MustBuildMutable().
			WithTxnContext(txnCtx).
			WithTimeTick(timetick).
			WithLastConfirmed(walimplstest.NewTestMessageID(1)).
			IntoImmutableMessage(walimplstest.NewTestMessageID(1)),
	}
	return immutables
}
