package impl

import (
	"context"
	"errors"
	"testing"

	"github.com/Masterminds/semver/v3"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	cluster_job_mock "github.com/tidbcloud/dataflow-service/app/internal/cluster_job/mock"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	backup_repo "github.com/tidbcloud/dataflow-service/app/repo/backup"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	repo_mock "github.com/tidbcloud/dataflow-service/app/repo/mock"
	internal_config "github.com/tidbcloud/dataflow-service/internal/config"
	bizerr_wrapper "github.com/tidbcloud/dataflow-service/pkg/bizerr-wrapper"
	grpcmock "github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli/mock"
	"github.com/tidbcloud/pkgv2/bizerr"
	dedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"gorm.io/gorm"
)

type pitBackupSwitcherFakeTransProvider struct {
	postCallbackErr      error
	callbackContextKey   interface{}
	callbackContextValue interface{}
}

func (f *pitBackupSwitcherFakeTransProvider) ReadDB(_ context.Context) *gorm.DB {
	return nil
}

func (f *pitBackupSwitcherFakeTransProvider) WriteDB(_ context.Context) *gorm.DB {
	return nil
}

func (f *pitBackupSwitcherFakeTransProvider) Transaction(
	ctx context.Context, cb func(ctx context.Context) error,
) error {
	txCtx := ctx
	if f.callbackContextKey != nil {
		txCtx = context.WithValue(txCtx, f.callbackContextKey, f.callbackContextValue)
	}
	if err := cb(txCtx); err != nil {
		return err
	}
	return f.postCallbackErr
}

type pitBackupSwitcherTxCtxKey struct{}

type pitBackupSwitcherTxState struct {
	staged map[uint64]entity.PitBackupSwitch
}

type pitBackupSwitcherInMemoryTransProvider struct {
	store map[uint64]entity.PitBackupSwitch
}

func newPitBackupSwitcherInMemoryTransProvider() *pitBackupSwitcherInMemoryTransProvider {
	return &pitBackupSwitcherInMemoryTransProvider{
		store: make(map[uint64]entity.PitBackupSwitch),
	}
}

func (p *pitBackupSwitcherInMemoryTransProvider) ReadDB(_ context.Context) *gorm.DB {
	return nil
}

func (p *pitBackupSwitcherInMemoryTransProvider) WriteDB(_ context.Context) *gorm.DB {
	return nil
}

func (p *pitBackupSwitcherInMemoryTransProvider) Transaction(
	ctx context.Context, cb func(ctx context.Context) error,
) error {
	staged := make(map[uint64]entity.PitBackupSwitch, len(p.store))
	for id, sw := range p.store {
		staged[id] = copyPitBackupSwitchEntity(sw)
	}
	txState := &pitBackupSwitcherTxState{staged: staged}
	txCtx := context.WithValue(ctx, pitBackupSwitcherTxCtxKey{}, txState)
	if err := cb(txCtx); err != nil {
		return err
	}
	p.store = txState.staged
	return nil
}

type pitBackupSwitcherInMemoryRepo struct {
	transProvider *pitBackupSwitcherInMemoryTransProvider
}

func (r *pitBackupSwitcherInMemoryRepo) Save(ctx context.Context, sw *entity.PitBackupSwitch) bizerr.BizError {
	return r.save(ctx, sw)
}

func (r *pitBackupSwitcherInMemoryRepo) save(ctx context.Context, sw *entity.PitBackupSwitch) bizerr.BizError {
	copied := copyPitBackupSwitchEntity(*sw)
	txState, inTx := ctx.Value(pitBackupSwitcherTxCtxKey{}).(*pitBackupSwitcherTxState)
	if inTx {
		txState.staged[copied.ID.Uint64()] = copied
		return nil
	}
	r.transProvider.store[copied.ID.Uint64()] = copied
	return nil
}

func (r *pitBackupSwitcherInMemoryRepo) Load(_ context.Context, id uint64) (*entity.PitBackupSwitch, bizerr.BizError) {
	sw, ok := r.transProvider.store[id]
	if !ok {
		return nil, bizerr_wrapper.NewNotFoundErr("pit backup switch %d not found", id)
	}
	copied := copyPitBackupSwitchEntity(sw)
	return &copied, nil
}

func (r *pitBackupSwitcherInMemoryRepo) LoadByClusterID(
	_ context.Context, clusterID uint64,
) (*entity.PitBackupSwitch, bizerr.BizError) {
	for _, sw := range r.transProvider.store {
		if sw.ClusterID.Uint64() == clusterID {
			copied := copyPitBackupSwitchEntity(sw)
			return &copied, nil
		}
	}
	return nil, bizerr_wrapper.NewNotFoundErr("pit backup switch for cluster %d not found", clusterID)
}

func (r *pitBackupSwitcherInMemoryRepo) List(
	ctx context.Context, _ ...backup_repo.PitBackupSwitchFilterOp,
) ([]entity.PitBackupSwitch, bizerr.BizError) {
	var source map[uint64]entity.PitBackupSwitch
	txState, inTx := ctx.Value(pitBackupSwitcherTxCtxKey{}).(*pitBackupSwitcherTxState)
	if inTx {
		source = txState.staged
	} else {
		source = r.transProvider.store
	}
	result := make([]entity.PitBackupSwitch, 0, len(source))
	for _, sw := range source {
		result = append(result, copyPitBackupSwitchEntity(sw))
	}
	return result, nil
}

func copyPitBackupSwitchEntity(sw entity.PitBackupSwitch) entity.PitBackupSwitch {
	copied := sw
	if sw.Action.CurrentActions != nil {
		copied.Action.CurrentActions = append([]string(nil), sw.Action.CurrentActions...)
	}
	return copied
}

func testDedicatedCluster(clusterID uint64) *mgmt.Cluster {
	return &mgmt.Cluster{
		ClusterId:      clusterID,
		ProjectId:      2001,
		OrganizationId: 3001,
		ClusterInfo: &mgmt.Cluster_DedicatedCluster{
			DedicatedCluster: &dedicatedv1.Cluster{
				State:   dedicatedv1.Cluster_AVAILABLE,
				Version: "999.0.0",
			},
		},
	}
}

func testPitBackupSwitch(
	switchID, clusterID uint64,
	status entity.PitBackupSwitchStatus,
	actions []string,
) *entity.PitBackupSwitch {
	return &entity.PitBackupSwitch{
		Updatable: entity.Updatable{
			Base: entity.Base{ID: entity.ID(switchID)},
		},
		TenantID:  entity.ID(3001),
		ProjectID: entity.ID(2001),
		ClusterID: entity.ID(clusterID),
		Status:    status,
		Action: entity.PitBackupSwitchAction{
			CurrentActions: actions,
		},
	}
}

func ensureEnableVersionConstraintForTest(t *testing.T) {
	t.Helper()

	oldConstraint := internal_config.Global().BackupConfig.LogBackupEnableMinVerConstraint
	constraint, err := semver.NewConstraint(">= 0.0.0")
	require.NoError(t, err)
	internal_config.Global().BackupConfig.LogBackupEnableMinVerConstraint = constraint

	t.Cleanup(func() {
		internal_config.Global().BackupConfig.LogBackupEnableMinVerConstraint = oldConstraint
	})
}

func TestEnable_RegistersPauseLockAfterSavingSwitch(t *testing.T) {
	ensureEnableVersionConstraintForTest(t)

	const (
		clusterID  uint64 = 1001
		switcherID uint64 = 2001
	)

	txCtxKey := struct{}{}
	transProvider := &pitBackupSwitcherFakeTransProvider{
		callbackContextKey:   txCtxKey,
		callbackContextValue: true,
	}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(switcherID, clusterID, entity.PitBackupStatusDisabled, nil)
	pitBackupSwitch.Turn = 41
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()

	saved := false
	registeredBeforeSave := false
	clusterJob.On("RegisterPiTRPauseLock", mock.Anything, clusterID, uint64(42)).
		Run(func(_ mock.Arguments) {
			if !saved {
				registeredBeforeSave = true
			}
		}).
		Return(nil).Once()
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Run(func(args mock.Arguments) {
			saved = true
			saveCtx, ok := args.Get(0).(context.Context)
			require.True(t, ok)
			assert.Equal(t, true, saveCtx.Value(txCtxKey))
			savedSwitch, ok := args.Get(1).(*entity.PitBackupSwitch)
			require.True(t, ok)
			assert.Equal(t, entity.PitBackupStatusWaitingForEnabling, savedSwitch.Status)
			assert.Equal(t, uint64(42), savedSwitch.Turn)
		}).
		Return(nil).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, err := switcher.Enable(context.Background(), &domain.EnablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, switcherID, resp.PitBackupSwitchID)
	assert.False(t, registeredBeforeSave)
	clusterJob.AssertNotCalled(t, "DeregisterPiTRPauseLock", mock.Anything, mock.Anything)
}

func TestCancelRestart_CompensatedWaitingStateWithEmptyActions(t *testing.T) {
	const (
		clusterID uint64 = 1010
		switchID  uint64 = 2010
	)

	transProvider := &pitBackupSwitcherFakeTransProvider{}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	pitBackupSwitch := testPitBackupSwitch(
		switchID,
		clusterID,
		entity.PitBackupStatusWaitingForRestarting,
		[]string{},
	)
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Run(func(args mock.Arguments) {
			savedSwitch, ok := args.Get(1).(*entity.PitBackupSwitch)
			require.True(t, ok)
			assert.Equal(t, entity.PitBackupStatusEnabled, savedSwitch.Status)
			assert.Empty(t, savedSwitch.Action.CurrentActions)
		}).
		Return(nil).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		nil,
		nil,
	)

	err := switcher.CancelRestart(context.Background(), &domain.CancelRestartPitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.NoError(t, err)
}

func TestEnable_TransactionFailureDoesNotRegisterPauseLock(t *testing.T) {
	ensureEnableVersionConstraintForTest(t)

	const (
		clusterID uint64 = 1002
		switchID  uint64 = 2002
	)

	transProvider := &pitBackupSwitcherFakeTransProvider{postCallbackErr: errors.New("commit failed")}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(switchID, clusterID, entity.PitBackupStatusDisabled, nil)
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Return(nil).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, err := switcher.Enable(context.Background(), &domain.EnablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.Error(t, err)
	assert.Nil(t, resp)
	clusterJob.AssertNotCalled(t, "RegisterPiTRPauseLock", mock.Anything, mock.Anything, mock.Anything)
	clusterJob.AssertNotCalled(t, "DeregisterPiTRPauseLock", mock.Anything, mock.Anything)
}

func TestEnable_RegisterPauseLockFailureDoesNotFailEnable(t *testing.T) {
	ensureEnableVersionConstraintForTest(t)

	const (
		clusterID uint64 = 1004
		switchID  uint64 = 2004
	)

	transProvider := &pitBackupSwitcherFakeTransProvider{}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(switchID, clusterID, entity.PitBackupStatusDisabled, nil)
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()
	clusterJob.On("RegisterPiTRPauseLock", mock.Anything, clusterID, uint64(1)).
		Return(bizerr_wrapper.NewInternalErr("register lock failed")).Once()
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Run(func(args mock.Arguments) {
			savedSwitch, ok := args.Get(1).(*entity.PitBackupSwitch)
			require.True(t, ok)
			assert.Equal(t, entity.PitBackupStatusWaitingForEnabling, savedSwitch.Status)
			assert.Equal(t, uint64(1), savedSwitch.Turn)
		}).
		Return(nil).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, err := switcher.Enable(context.Background(), &domain.EnablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, switchID, resp.PitBackupSwitchID)
	clusterJob.AssertNotCalled(t, "DeregisterPiTRPauseLock", mock.Anything, mock.Anything)
}

func TestDisableWaitingForEnabling_DeregistersPauseLockBeforeSavingSwitch(t *testing.T) {
	const (
		clusterID uint64 = 1003
		switchID  uint64 = 2003
	)

	transProvider := &pitBackupSwitcherFakeTransProvider{}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(switchID, clusterID, entity.PitBackupStatusWaitingForEnabling, nil)
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()

	deregistered := false
	deregisteredBeforeSave := false
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Run(func(_ mock.Arguments) {
			deregisteredBeforeSave = deregistered
		}).
		Return(nil).Once()
	clusterJob.On("DeregisterPiTRPauseLock", mock.Anything, clusterID, uint64(0)).
		Run(func(_ mock.Arguments) {
			deregistered = true
		}).
		Return(nil).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, err := switcher.Disable(context.Background(), &domain.DisablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, switchID, resp.PitBackupSwitchID)
	assert.True(t, deregisteredBeforeSave)
}

func TestDisableWaitingForEnabling_CommitFailureStillDeregistersPauseLock(t *testing.T) {
	const (
		clusterID uint64 = 1006
		switchID  uint64 = 2006
	)

	transProvider := &pitBackupSwitcherFakeTransProvider{postCallbackErr: errors.New("commit failed")}
	pitBackupSwitchRepo := repo_mock.NewPitBackupSwitchRepo(t)
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(
		switchID,
		clusterID,
		entity.PitBackupStatusWaitingForEnabling,
		[]string{string(entity.PitBackupSwitchActionEnable)},
	)
	pitBackupSwitchRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.PitBackupSwitch{*pitBackupSwitch}, nil).Once()
	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()
	pitBackupSwitchRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.PitBackupSwitch")).
		Return(nil).Once()

	deregistered := false
	clusterJob.On("DeregisterPiTRPauseLock", mock.Anything, clusterID, uint64(0)).
		Run(func(_ mock.Arguments) { deregistered = true }).
		Return(nil).Once()
	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, bizErr := switcher.Disable(context.Background(), &domain.DisablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.Error(t, bizErr)
	assert.Nil(t, resp)
	assert.True(t, deregistered)
	clusterJob.AssertNotCalled(t, "RegisterPiTRPauseLock", mock.Anything, mock.Anything)
}

func TestDisableWaitingForEnabling_DeregisterPauseLockFailure(t *testing.T) {
	const (
		clusterID uint64 = 1005
		switchID  uint64 = 2005
	)

	transProvider := newPitBackupSwitcherInMemoryTransProvider()
	pitBackupSwitchRepo := &pitBackupSwitcherInMemoryRepo{transProvider: transProvider}
	logBackupRepo := repo_mock.NewLogBackupRepo(t)
	mgmtCli := grpcmock.NewMgmtSvcGrpcClient(t)
	clusterJob := cluster_job_mock.NewModel(t)

	pitBackupSwitch := testPitBackupSwitch(
		switchID,
		clusterID,
		entity.PitBackupStatusWaitingForEnabling,
		[]string{string(entity.PitBackupSwitchActionEnable)},
	)
	require.NoError(t, pitBackupSwitchRepo.save(context.Background(), pitBackupSwitch))

	mgmtCli.On("GetAnyDedicatedClusterBasicByID", mock.Anything, clusterID).
		Return(testDedicatedCluster(clusterID), nil).Once()
	clusterJob.On("DeregisterPiTRPauseLock", mock.Anything, clusterID, uint64(0)).
		Return(bizerr_wrapper.NewInternalErr("deregister lock failed")).Once()

	switcher := NewPitBackupSwitcher(
		transProvider,
		pitBackupSwitchRepo,
		logBackupRepo,
		nil,
		mgmtCli,
		clusterJob,
	)

	resp, bizErr := switcher.Disable(context.Background(), &domain.DisablePitBackupReq{
		SwitchPitBackupReq: domain.SwitchPitBackupReq{ClusterID: clusterID},
	})
	require.Error(t, bizErr)
	assert.Nil(t, resp)

	reloaded, loadErr := pitBackupSwitchRepo.Load(context.Background(), pitBackupSwitch.ID.Uint64())
	require.NoError(t, loadErr)
	assert.Equal(t, entity.PitBackupStatusWaitingForEnabling, reloaded.Status)
}