package server

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

	"github.com/pingcap/kvproto/pkg/metapb"
	"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"
	braction "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	clustergateway_mocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	dbmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/dbmocks"
	executor_mocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	repoext_mocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repoext/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gorm.io/gorm"
)

func TestServerImpl_GetRestoreResource(t *testing.T) {
	tests := []struct {
		name            string
		req             *br.CreateLightweightRestoreInput
		setupMocks      func(*repomocks.SharedPoolRepo, *dataplanemocks.Manager)
		bizCfg          *config.BizConfig
		expectedCpu     string
		expectedMemory  string
		expectedStorage string
		expectedError   bool
		errorContains   string
	}{
		{
			name: "all resources provided in request",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "4",
					ResourceMemory:     "16Gi",
					ResourceStorage:    "200Gi",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				// No mocks needed since all resources are provided
			},
			expectedCpu:     "4",
			expectedMemory:  "16Gi",
			expectedStorage: "200Gi",
			expectedError:   false,
		},
		{
			name: "resources with whitespace trimmed",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "  8  ",
					ResourceMemory:     "  32Gi  ",
					ResourceStorage:    "  500Gi  ",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				// No mocks needed since all resources are provided
			},
			expectedCpu:     "8",
			expectedMemory:  "32Gi",
			expectedStorage: "500Gi",
			expectedError:   false,
		},
		{
			name: "only cpu and memory provided - should get storage from config",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "4",
					ResourceMemory:     "16Gi",
					ResourceStorage:    "",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				// No mocks needed since cpu and memory are provided
			},
			bizCfg: &config.BizConfig{
				DefaultTiBRConfig: config.TiBRConfig{
					StorageSize: "150Gi",
				},
			},
			expectedCpu:     "4",
			expectedMemory:  "16Gi",
			expectedStorage: "150Gi",
			expectedError:   false,
		},
		{
			name: "only cpu and memory provided - should get storage from default value",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "4",
					ResourceMemory:     "16Gi",
					ResourceStorage:    "",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				// No mocks needed since cpu and memory are provided
			},
			expectedCpu:     "4",
			expectedMemory:  "16Gi",
			expectedStorage: braction.DefaultTiBRStorageSize,
			expectedError:   false,
		},
		{
			name: "shared pool load error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "",
					ResourceMemory:     "",
					ResourceStorage:    "",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				sharedPoolRepo.On("Load", context.Background(), "pool-123", false).Return(nil, errors.New("database connection failed"))
			},
			expectedError: true,
			errorContains: "failed to load sharedPool: database connection failed",
		},
		{
			name: "dataplane manager get client error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					ResourceCpu:        "",
					ResourceMemory:     "",
					ResourceStorage:    "",
					TargetSharedPoolId: "pool-123",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						DataplaneName: "test-dataplane",
					},
				}
				sharedPoolRepo.On("Load", context.Background(), "pool-123", false).Return(sharedPool, nil)
				dataplaneManager.On("GetClient", "test-dataplane").Return(nil, errors.New("failed to connect to dataplane"))
			},
			expectedError: true,
			errorContains: "failed to get remote client for shared pool: failed to connect to dataplane",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
			dataplaneManager := dataplanemocks.NewManager(t)

			// Setup mock expectations
			tt.setupMocks(sharedPoolRepo, dataplaneManager)

			// Create server
			server := &serverImpl{
				sharedPoolRepo:   sharedPoolRepo,
				dataplaneManager: dataplaneManager,
				bizCfg:           tt.bizCfg,
			}

			// Call the method
			ctx := context.Background()
			cpu, memory, storage, err := server.getRestoreResource(ctx, tt.req)

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

			// Check success expectations
			if err != nil {
				t.Errorf("unexpected error: %v", err)
				return
			}

			if cpu != tt.expectedCpu {
				t.Errorf("expected cpu %q, got %q", tt.expectedCpu, cpu)
			}

			if memory != tt.expectedMemory {
				t.Errorf("expected memory %q, got %q", tt.expectedMemory, memory)
			}

			if storage != tt.expectedStorage {
				t.Errorf("expected storage %q, got %q", tt.expectedStorage, storage)
			}
		})
	}
}

func TestServerImpl_CreateLightweightRestore(t *testing.T) {
	tests := []struct {
		name          string
		req           *br.CreateLightweightRestoreInput
		setupMocks    func(*repomocks.SharedPoolRepo, *repomocks.LightweightRestoreRepo, *repoext_mocks.KeyspaceFetcher, *clustergateway_mocks.Client, *executor_mocks.RemoteExecutor, *dataplanemocks.Manager, *dbmocks.Transactioner)
		expectedError bool
		errorContains string
		description   string
	}{
		{
			name: "successful restore creation",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:           "restore-123",
					SourceClusterId:     "source-cluster-456",
					TargetClusterId:     "target-cluster-789",
					SourceSharedPoolId:  "source-pool-111",
					TargetSharedPoolId:  "target-pool-222",
					ResourceCpu:         "4",
					ResourceMemory:      "16Gi",
					ResourceStorage:     "200Gi",
					KernelBackupId:      12345,
					KernelBackupName:    "backup-name",
					BackupDataSizeBytes: 1024000,
					PointInTime:         "2023-12-01T10:00:00Z",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				// Mock keyspace fetching
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)

				// Mock shared pool loading
				targetSharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:           "target-pool-222",
						CapacityMode: apiv1.CapacityMode_RCU_DRIVEN,
					},
				}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)

				// Mock cluster gateway
				targetPdCluster := &metapb.Cluster{
					Id: 98765,
				}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)

				// Mock restore record check (should return nil for new restore)
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(nil, status.Error(codes.NotFound, "restore not found"))

				// Mock DAG execution
				dagExecutor.On("Submit", mock.Anything, dag.PlanCreateLightweightXRestore, mock.Anything, "create-lw-x-restore-restore-123").Return(nil)

				// Mock restore record creation
				lwRestoreRepo.On("Create", context.Background(), mock.AnythingOfType("*entity.LightweightRestore")).Return(nil)

				// Mock successful transaction that executes the callback
				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)
			},
			expectedError: false,
			description:   "Should successfully create a new lightweight restore",
		},
		{
			name: "target keyspace fetch error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(&entity.SharedPool{}, nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("", errors.New("keyspace not found"))
			},
			expectedError: true,
			errorContains: "get targetKeyspace for cluster target-cluster-789 failed",
			description:   "Should fail when target keyspace cannot be fetched",
		},
		{
			name: "source keyspace fetch error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(&entity.SharedPool{}, nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("", errors.New("source keyspace not found"))
			},
			expectedError: true,
			errorContains: "get sourceKeyspace for cluster source-cluster-456 failed",
			description:   "Should fail when source keyspace cannot be fetched",
		},
		{
			name: "shared pool load error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(nil, errors.New("shared pool not found"))
			},
			expectedError: true,
			errorContains: "failed to load targetSharedPool",
			description:   "Should fail when target shared pool cannot be loaded",
		},
		{
			name: "get cluster error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(nil, errors.New("cluster not found"))
			},
			expectedError: true,
			errorContains: "failed to get target pd cluster id",
			description:   "Should fail when target PD cluster cannot be retrieved",
		},
		{
			name: "restore record load error",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(nil, errors.New("database error"))
			},
			expectedError: true,
			errorContains: "failed to load lw restore",
			description:   "Should fail when restore record check fails",
		},
		{
			name: "restore already exists - active restore",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)

				existingRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID: "restore-123",
					},
				}
				// Simulate not deleted (DeletedAt not set)
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(existingRestore, nil)
			},
			expectedError: true,
			errorContains: "restore restore-123 already exists, please don't create it again",
			description:   "Should fail when restore already exists and is active",
		},
		{
			name: "restore already exists - deleted restore conflict",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)

				existingRestore := &entity.LightweightRestore{
					LightweightRestore: schema.LightweightRestore{
						RestoreID: "restore-123",
					},
				}
				// Simulate deleted restore (set DeletedAt manually in test)
				existingRestore.DeletedAt = gorm.DeletedAt{Valid: true}
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(existingRestore, nil)
			},
			expectedError: true,
			errorContains: "restore id conflict with a deleted restore, please use a new restore id",
			description:   "Should fail when restore ID conflicts with a deleted restore",
		},
		{
			name: "DAG execution failure",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
					ResourceCpu:        "4",
					ResourceMemory:     "16Gi",
					ResourceStorage:    "200Gi",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				// Setup successful mocks until DAG execution
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(nil, nil)

				// DAG execution fails
				dagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("DAG execution failed"))

				// Mock transaction that executes the 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 the mocked errors above
					if err == nil {
						t.Errorf("expected error but got none")
					}
				}).Return(errors.New("transaction failed due to inner error"))
			},
			expectedError: true,
			errorContains: "failed to create lightweight restore",
			description:   "Should fail when DAG execution fails",
		},
		{
			name: "restore record creation failure",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:          "restore-123",
					SourceClusterId:    "source-cluster-456",
					TargetClusterId:    "target-cluster-789",
					SourceSharedPoolId: "source-pool-111",
					TargetSharedPoolId: "target-pool-222",
					ResourceCpu:        "4",
					ResourceMemory:     "16Gi",
					ResourceStorage:    "200Gi",
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				// Setup successful mocks until restore record creation
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-222").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-111").Return("source-keyspace", nil)
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
				sharedPoolRepo.On("Load", context.Background(), "target-pool-222", false).Return(targetSharedPool, nil)
				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)
				lwRestoreRepo.On("Load", context.Background(), "restore-123", true).Return(nil, nil)

				// DAG execution succeeds but restore record creation fails
				dagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
				lwRestoreRepo.On("Create", context.Background(), mock.AnythingOfType("*entity.LightweightRestore")).Return(errors.New("database insert failed"))

				// Mock Abort for rollback after transaction failure
				dagExecutor.On("Abort", mock.Anything, mock.Anything).Return(nil)

				// Mock transaction that executes the 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 the mocked errors above
					if err == nil {
						t.Errorf("expected error but got none")
					}
				}).Return(errors.New("transaction failed due to inner error"))
			},
			expectedError: true,
			errorContains: "failed to create lightweight restore",
			description:   "Should fail when restore record creation fails",
		},
		{
			name: "topology driven pools use fetched keyspace",
			req: &br.CreateLightweightRestoreInput{
				CreateLightweightRestoreRequest: &apiv1.CreateLightweightRestoreRequest{
					RestoreId:           "restore-topology-123",
					SourceClusterId:     "source-cluster-456",
					TargetClusterId:     "target-cluster-789",
					SourceSharedPoolId:  "source-pool-topology",
					TargetSharedPoolId:  "target-pool-topology",
					ResourceCpu:         "4",
					ResourceMemory:      "16Gi",
					ResourceStorage:     "200Gi",
					KernelBackupName:    "backup-name",
					BackupDataSizeBytes: 1024,
				},
			},
			setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, lwRestoreRepo *repomocks.LightweightRestoreRepo, keyspaceFetcher *repoext_mocks.KeyspaceFetcher, clustergateway *clustergateway_mocks.Client, dagExecutor *executor_mocks.RemoteExecutor, dataplaneManager *dataplanemocks.Manager, transaction *dbmocks.Transactioner) {
				targetSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "target-pool-topology", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}
				keyspaceFetcher.On("GetKeyspace", context.Background(), "target-cluster-789", "target-pool-topology").Return("target-keyspace", nil)
				keyspaceFetcher.On("GetKeyspace", context.Background(), "source-cluster-456", "source-pool-topology").Return("source-keyspace", nil)
				sharedPoolRepo.On("Load", context.Background(), "target-pool-topology", false).Return(targetSharedPool, nil)

				targetPdCluster := &metapb.Cluster{Id: 98765}
				clustergateway.On("GetCluster", context.Background(), targetSharedPool).Return(targetPdCluster, nil)

				lwRestoreRepo.On("Load", context.Background(), "restore-topology-123", true).Return(nil, status.Error(codes.NotFound, "restore not found"))
				dagExecutor.On("Submit", mock.Anything, dag.PlanCreateLightweightXRestore, mock.Anything, "create-lw-x-restore-restore-topology-123").Return(nil)
				lwRestoreRepo.On("Create", context.Background(), mock.MatchedBy(func(r *entity.LightweightRestore) bool {
					return r != nil && r.TargetKeyspace == "target-keyspace" && r.SourceKeyspace == "source-keyspace"
				})).Return(nil)

				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())
					if err != nil {
						t.Errorf("failed to execute callback: %v", err)
					}
				}).Return(nil)
			},
			expectedError: false,
			description:   "Should resolve keyspace through keyspace fetcher for topology-driven pools",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
			lwRestoreRepo := repomocks.NewLightweightRestoreRepo(t)
			keyspaceFetcher := repoext_mocks.NewKeyspaceFetcher(t)
			clustergateway := clustergateway_mocks.NewClient(t)
			dagExecutor := executor_mocks.NewRemoteExecutor(t)
			dataplaneManager := dataplanemocks.NewManager(t)
			transaction := dbmocks.NewTransactioner(t)

			// Setup mock expectations
			tt.setupMocks(sharedPoolRepo, lwRestoreRepo, keyspaceFetcher, clustergateway, dagExecutor, dataplaneManager, transaction)

			// Create server
			server := &serverImpl{
				sharedPoolRepo:    sharedPoolRepo,
				lwRestoreRepo:     lwRestoreRepo,
				keyspaceFetcher:   keyspaceFetcher,
				clustergatewayCli: clustergateway,
				dagExecutor:       dagExecutor,
				dataplaneManager:  dataplaneManager,
				transactioner:     transaction,
				bizCfg: &config.BizConfig{
					DefaultTiBRConfig: config.TiBRConfig{
						StorageSize: "100Gi",
					},
				},
			}

			// Call the method
			ctx := context.Background()
			err := server.CreateLightweightRestore(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)
			}
		})
	}
}