package server

import (
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/stretchr/testify/mock"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/config"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/server/br"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	dbmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/dbmocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor"
	executor_mocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executorext"
	executorext_mocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executorext/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"gorm.io/gorm"
)

func TestServerImpl_DeleteLightweightRestore(t *testing.T) {
	tests := []struct {
		name          string
		req           *br.DeleteLightweightRestoreInput
		setupMocks    func(*repomocks.LightweightRestoreRepo, *executor_mocks.RemoteExecutor, *executorext_mocks.DagInstanceFinder, *dbmocks.Transactioner)
		expectedError bool
		errorContains string
		description   string
	}{
		{
			name: "restore already deleted - should return success",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock restore that is already deleted
				deletedRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
					},
				}
				deletedRestore.DeletedAt = gorm.DeletedAt{Valid: true}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(deletedRestore, nil)
			},
			expectedError: false,
			description:   "Should return success immediately if restore is already deleted",
		},
		{
			name: "restore load error",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(nil, errors.New("database connection failed"))
			},
			expectedError: true,
			errorContains: "failed to load restore",
			description:   "Should fail when restore cannot be loaded",
		},
		{
			name: "successful deletion - no running dags",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock no running creation dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock no running deletion dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanDeleteLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock successful deletion dag creation
				dagExecutor.On("Submit", mock.Anything, dag.PlanDeleteLightweightXRestore, mock.Anything, "delete-lw-x-restore-restore-123").Return(nil)
			},
			expectedError: false,
			description:   "Should successfully create deletion dag when no existing dags are running",
		},
		{
			name: "abort running creation dag before deletion",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock running creation dag
				runningCreationDag := &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							ExecId: "exec-123",
						},
					},
				}
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(runningCreationDag, nil)

				// Mock successful abortion of creation dag
				dagExecutor.On("Abort", mock.Anything, "exec-123").Return(nil)
				lwRestoreRepo.On("SetStatus", context.Background(), "restore-123", schema.LightweightRestoreStatusCancelled).Return(nil)

				// Mock successful transaction for abortion
				transaction.On("Transaction", context.Background(), mock.AnythingOfType("func(context.Context) error")).Run(func(args mock.Arguments) {
					callback := args.Get(1).(func(context.Context) error)
					err := callback(context.Background()) // Execute the callback successfully
					if err != nil {
						t.Errorf("failed to execute callback: %v", err)
					}
				}).Return(nil)

				// Mock no running deletion dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanDeleteLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock successful deletion dag creation
				dagExecutor.On("Submit", mock.Anything, dag.PlanDeleteLightweightXRestore, mock.Anything, "delete-lw-x-restore-restore-123").Return(nil)
			},
			expectedError: false,
			description:   "Should abort running creation dag before creating deletion dag",
		},
		{
			name: "running deletion dag already exists",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock no running creation dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock running deletion dag
				runningDeletionDag := &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							ExecId: "exec-delete-123",
						},
					},
				}
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanDeleteLightweightXRestore, mock.Anything).Return(runningDeletionDag, nil)
			},
			expectedError: false,
			description:   "Should return success when deletion dag is already running",
		},
		{
			name: "dag instance finder error for creation dag",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock error when finding creation dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(nil, errors.New("dag finder error"))
			},
			expectedError: true,
			errorContains: "dag finder error",
			description:   "Should fail when dag instance finder returns error",
		},
		{
			name: "abort creation dag fails",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock running creation dag
				runningCreationDag := &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							ExecId: "exec-123",
						},
					},
				}
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(runningCreationDag, nil)

				// Mock failed abortion of creation dag
				dagExecutor.On("Abort", mock.Anything, "exec-123").Return(errors.New("abort failed"))

				// Mock transaction that executes callback (which will fail)
				transaction.On("Transaction", context.Background(), mock.AnythingOfType("func(context.Context) error")).Run(func(args mock.Arguments) {
					callback := args.Get(1).(func(context.Context) error)
					err := callback(context.Background()) // This will fail due to abort error
					if err == nil {
						t.Errorf("expected error but got none")
					}
				}).Return(errors.New("transaction failed"))
			},
			expectedError: true,
			errorContains: "abort existed",
			description:   "Should fail when aborting creation dag fails",
		},
		{
			name: "set status to cancelled fails",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock running creation dag
				runningCreationDag := &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							ExecId: "exec-123",
						},
					},
				}
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(runningCreationDag, nil)

				// Mock successful abortion but failed status update
				dagExecutor.On("Abort", mock.Anything, "exec-123").Return(nil)
				lwRestoreRepo.On("SetStatus", context.Background(), "restore-123", schema.LightweightRestoreStatusCancelled).Return(errors.New("database update failed"))

				// Mock transaction that executes callback (which will fail)
				transaction.On("Transaction", context.Background(), mock.AnythingOfType("func(context.Context) error")).Run(func(args mock.Arguments) {
					callback := args.Get(1).(func(context.Context) error)
					err := callback(context.Background()) // This will fail due to SetStatus error
					if err == nil {
						t.Errorf("expected error but got none")
					}
				}).Return(errors.New("transaction failed"))
			},
			expectedError: true,
			errorContains: "abort existed",
			description:   "Should fail when setting status to cancelled fails",
		},
		{
			name: "deletion dag creation fails",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock no running creation dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock no running deletion dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanDeleteLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock failed deletion dag creation
				dagExecutor.On("Submit", mock.Anything, dag.PlanDeleteLightweightXRestore, mock.Anything, "delete-lw-x-restore-restore-123").Return(errors.New("dag creation failed"))
			},
			expectedError: true,
			errorContains: "failed to submit",
			description:   "Should fail when deletion dag creation fails",
		},
		{
			name: "dag instance finder error for deletion dag",
			req: &br.DeleteLightweightRestoreInput{
				DeleteLightweightRestoreRequest: &apiv1.DeleteLightweightRestoreRequest{
					RestoreId:          "restore-123",
					TargetSharedPoolId: "pool-456",
				},
			},
			setupMocks: func(lwRestoreRepo *repomocks.LightweightRestoreRepo, dagExecutor *executor_mocks.RemoteExecutor, dagInstanceFinder *executorext_mocks.DagInstanceFinder, transaction *dbmocks.Transactioner) {
				// Mock active restore
				activeRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID:          "restore-123",
						TargetSharedPoolID: "pool-456",
						SourceSharedPoolID: "pool-111",
					},
				}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(activeRestore, nil)

				// Mock no running creation dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanCreateLightweightXRestore, mock.Anything).Return(nil, executorext.DagInsNotFoundErr)

				// Mock error when finding deletion dag
				dagInstanceFinder.On("GetRunningDagInstance", context.Background(), dag.PlanDeleteLightweightXRestore, mock.Anything).Return(nil, errors.New("deletion dag finder error"))
			},
			expectedError: true,
			errorContains: "deletion dag finder error",
			description:   "Should fail when dag instance finder returns error for deletion dag",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			lwRestoreRepo := repomocks.NewLightweightRestoreRepo(t)
			dagExecutor := executor_mocks.NewRemoteExecutor(t)
			dagInstanceFinder := executorext_mocks.NewDagInstanceFinder(t)
			transaction := dbmocks.NewTransactioner(t)

			// Setup mock expectations
			tt.setupMocks(lwRestoreRepo, dagExecutor, dagInstanceFinder, transaction)

			// Create server
			server := &serverImpl{
				lwRestoreRepo:     lwRestoreRepo,
				sharedPoolRepo:    availableSharedPoolRepoForMutationTest(),
				dagExecutor:       dagExecutor,
				dagInstanceFinder: dagInstanceFinder,
				transactioner:     transaction,
				bizCfg: &config.BizConfig{
					DefaultTiBRConfig: config.TiBRConfig{
						StorageSize: "100Gi",
					},
				},
			}

			// Call the method
			ctx := context.Background()
			err := server.DeleteLightweightRestore(ctx, tt.req)

			// Check expectations
			if tt.expectedError {
				if err == nil {
					t.Errorf("expected error but got none - %s", tt.description)
					return
				}
				if tt.errorContains != "" {
					if !strings.Contains(err.Error(), tt.errorContains) {
						t.Errorf("expected error to contain %q, got %q - %s", tt.errorContains, err.Error(), tt.description)
					}
				}
				return
			}

			if err != nil {
				t.Errorf("unexpected error: %v - %s", err, tt.description)
			}
		})
	}
}