package br

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	braction "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br"
	"github.com/tidbcloud/cluster-service-ng/pkg/clustergateway"
	clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
)

func TestLwRestoreCrossSharedPoolAdaptor_GetRestorePackedBackupOutput(t *testing.T) {
	tests := []struct {
		name                  string
		restore               *entity.LightweightRestore
		targetSharedPool      *entity.SharedPool
		setupMocks            func(*clustergatewaymocks.Client)
		expectedStatus        apiv1.LightweightRestoreStatus
		expectedMessage       string
		expectedError         bool
		expectedErrorContains string
		expectedProgress      uint32
		expectedBytes         uint64
	}{
		{
			name: "task not found - wait for creation",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(nil, clustergateway.ErrBRTaskNotFound)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "wait for the kernel restore packed task created",
			expectedError:    false,
			expectedProgress: 50,
		},
		{
			name: "get task failed with other error",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(nil, errors.New("cluster connection failed"))
			},
			expectedError:         true,
			expectedErrorContains: "get kernel restore packed task failed: cluster connection failed",
		},
		{
			name: "task pending state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStatePending,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel is pending for restore packed",
			expectedError:    false,
			expectedProgress: 50,
		},
		{
			name: "task init state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateInit,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "wait for the kernel restore packed task initialized",
			expectedError:    false,
			expectedProgress: 60,
		},
		{
			name: "task running state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateRunning,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel is running for restore packed",
			expectedError:    false,
			expectedProgress: 90,
		},
		{
			name: "task done state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateDone,
					Info: map[string]any{
						"restore_bytes": uint64(2048),
					},
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel restore packed task done",
			expectedError:    false,
			expectedProgress: 100,
			expectedBytes:    2048,
		},
		{
			name: "task done state without restore_bytes",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateDone,
					Info:  map[string]any{},
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel restore packed task done",
			expectedError:    false,
			expectedProgress: 100,
			expectedBytes:    0,
		},
		{
			name: "task error state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateError,
					Error: "restore operation failed",
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "kernel restore packed task error: restore operation failed",
			expectedError:   false,
		},
		{
			name: "task error state with empty error message",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateError,
					Error: "",
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "kernel restore packed task error: ",
			expectedError:   false,
		},
		{
			name: "unknown task state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			targetSharedPool: &entity.SharedPool{},
			setupMocks: func(clustergatewayClient *clustergatewaymocks.Client) {
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskState("Unknown"),
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelRestorePackedTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "unknown state for kernel restore packed task",
			expectedError:   false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			clustergatewayClient := clustergatewaymocks.NewClient(t)

			// Setup mock expectations
			tt.setupMocks(clustergatewayClient)

			// Create adaptor
			adaptor := &lwRestoreCrossSharedPoolAdaptor{
				clustergatewayCli: clustergatewayClient,
				sharedPoolRepo:    mocks.NewSharedPoolRepo(t), // Not used in this method but required for struct
			}

			// Call the method
			ctx := context.Background()
			result, err := adaptor.getRestorePackedBackupOutput(ctx, tt.restore, tt.targetSharedPool)

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

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

			if result == nil {
				t.Errorf("expected result but got nil")
				return
			}

			if result.Status != tt.expectedStatus {
				t.Errorf("expected status %v, got %v", tt.expectedStatus, result.Status)
			}

			if result.Message != tt.expectedMessage {
				t.Errorf("expected message %q, got %q", tt.expectedMessage, result.Message)
			}

			if result.DataRestoreProgressPercentage != tt.expectedProgress {
				t.Errorf("expected progress %v, got %v", tt.expectedProgress, result.DataRestoreProgressPercentage)
			}

			if result.DataRestoreBytes != tt.expectedBytes {
				t.Errorf("expected bytes %v, got %v", tt.expectedBytes, result.DataRestoreBytes)
			}
		})
	}
}

func TestLwRestoreCrossSharedPoolAdaptor_GetPostActions(t *testing.T) {
	ctx := context.Background()
	sharedPoolRepo := mocks.NewSharedPoolRepo(t)
	globalSortPath := "s3://target-bucket/global-sort"
	sharedPoolRepo.On("Load", ctx, "target-pool", false).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS},
			ObjectStorage: &apiv1.SharedPool_ObjectStorage{
				Bucket: "target-bucket",
				Prefix: "",
			},
		},
	}, nil)

	adaptor := &lwRestoreCrossSharedPoolAdaptor{sharedPoolRepo: sharedPoolRepo}
	resp, err := adaptor.GetPostActions(ctx, &entity.LightweightRestore{
		LightweightRestore: schema.LightweightRestore{TargetSharedPoolID: "target-pool"},
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, "SET GLOBAL tidb_cloud_storage_uri = '"+globalSortPath+"'", resp.ExecSql)
	require.Equal(t, globalSortPath, resp.CloudStorageUri)
}

func TestLwRestoreCrossSharedPoolAdaptor_GetDataRestoringOutput(t *testing.T) {
	tests := []struct {
		name                  string
		restore               *entity.LightweightRestore
		setupMocks            func(*mocks.SharedPoolRepo, *clustergatewaymocks.Client)
		expectedStatus        apiv1.LightweightRestoreStatus
		expectedMessage       string
		expectedError         bool
		expectedErrorContains string
		expectedProgress      uint32
	}{
		{
			name: "shared pool load error",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(nil, errors.New("database connection failed"))
			},
			expectedError:         true,
			expectedErrorContains: "get target sharedPool pool-456 failed: database connection failed",
		},
		{
			name: "task not found - wait for creation",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(nil, clustergateway.ErrBRTaskNotFound)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "wait for the kernel pack backup task created",
			expectedError:    false,
			expectedProgress: 0,
		},
		{
			name: "get task failed with other error",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(nil, errors.New("cluster connection failed"))
			},
			expectedError:         true,
			expectedErrorContains: "get kernel pack backup task failed: cluster connection failed",
		},
		{
			name: "task pending state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStatePending,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel is pending for pack backup",
			expectedError:    false,
			expectedProgress: 0,
		},
		{
			name: "task init state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateInit,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "wait for the kernel pack backup task initialized",
			expectedError:    false,
			expectedProgress: 10,
		},
		{
			name: "task running state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateRunning,
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:   apiv1.LightweightRestoreStatus_DATA_RESTORING,
			expectedMessage:  "kernel is running for pack backup",
			expectedError:    false,
			expectedProgress: 40,
		},
		{
			name: "task error state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateError,
					Error: "backup operation failed",
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "kernel pack backup task error: backup operation failed",
			expectedError:   false,
		},
		{
			name: "task error state with empty error message",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskStateError,
					Error: "",
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "kernel pack backup task error: ",
			expectedError:   false,
		},
		{
			name: "unknown task state",
			restore: &entity.LightweightRestore{
				LightweightRestore: schema.LightweightRestore{
					RestoreID:          "restore-123",
					TargetSharedPoolID: "pool-456",
				},
			},
			setupMocks: func(sharedPoolRepo *mocks.SharedPoolRepo, clustergatewayClient *clustergatewaymocks.Client) {
				sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(&entity.SharedPool{}, nil)
				taskResponse := &clustergateway.TaskResponse{
					Id:    123,
					State: clustergateway.TaskState("Unknown"),
				}
				clustergatewayClient.On("GetTiBRTask", context.Background(),
					&entity.SharedPool{}, tioperator.TiBRCRNameOfRestore("restore-123"),
					braction.KernelPackBackupTaskId).Return(taskResponse, nil)
			},
			expectedStatus:  apiv1.LightweightRestoreStatus_DATA_RESTORING_ERROR,
			expectedMessage: "unknown state for kernel pack backup task",
			expectedError:   false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			sharedPoolRepo := mocks.NewSharedPoolRepo(t)
			clustergatewayClient := clustergatewaymocks.NewClient(t)

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

			// Create adaptor
			adaptor := &lwRestoreCrossSharedPoolAdaptor{
				clustergatewayCli: clustergatewayClient,
				sharedPoolRepo:    sharedPoolRepo,
			}

			// Call the method
			ctx := context.Background()
			result, err := adaptor.getDataRestoringOutput(ctx, tt.restore)

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

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

			if result == nil {
				t.Errorf("expected result but got nil")
				return
			}

			if result.Status != tt.expectedStatus {
				t.Errorf("expected status %v, got %v", tt.expectedStatus, result.Status)
			}

			if result.Message != tt.expectedMessage {
				t.Errorf("expected message %q, got %q", tt.expectedMessage, result.Message)
			}

			if result.DataRestoreProgressPercentage != tt.expectedProgress {
				t.Errorf("expected progress %v, got %v", tt.expectedProgress, result.DataRestoreProgressPercentage)
			}
		})
	}
}

func TestLwRestoreCrossSharedPoolAdaptor_GetDataRestoringOutput_TaskDoneState(t *testing.T) {
	// This test specifically handles the TaskDone state which calls getRestorePackedBackupOutput
	restore := &entity.LightweightRestore{
		LightweightRestore: schema.LightweightRestore{
			RestoreID:          "restore-123",
			TargetSharedPoolID: "pool-456",
		},
	}

	// Create mocks
	sharedPoolRepo := mocks.NewSharedPoolRepo(t)
	clustergatewayClient := clustergatewaymocks.NewClient(t)

	// Setup shared pool load
	targetSharedPool := &entity.SharedPool{}
	sharedPoolRepo.On("Load", context.Background(), "pool-456", false).Return(targetSharedPool, nil)

	// Setup pack backup task response - Done state
	packTaskResponse := &clustergateway.TaskResponse{
		Id:    123,
		State: clustergateway.TaskStateDone,
	}
	clustergatewayClient.On("GetTiBRTask", context.Background(),
		targetSharedPool, tioperator.TiBRCRNameOfRestore("restore-123"),
		braction.KernelPackBackupTaskId).Return(packTaskResponse, nil)

	// Setup restore packed task response - also Done state with restore bytes
	restoreTaskResponse := &clustergateway.TaskResponse{
		Id:    456,
		State: clustergateway.TaskStateDone,
		Info: map[string]any{
			"restore_bytes": uint64(1024),
		},
	}
	clustergatewayClient.On("GetTiBRTask", context.Background(),
		targetSharedPool, tioperator.TiBRCRNameOfRestore("restore-123"),
		braction.KernelRestorePackedTaskId).Return(restoreTaskResponse, nil)

	// Create adaptor
	adaptor := &lwRestoreCrossSharedPoolAdaptor{
		clustergatewayCli: clustergatewayClient,
		sharedPoolRepo:    sharedPoolRepo,
	}

	// Call the method
	ctx := context.Background()
	result, err := adaptor.getDataRestoringOutput(ctx, restore)

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

	if result == nil {
		t.Errorf("expected result but got nil")
		return
	}

	// When pack backup is done and restore packed is also done
	expectedStatus := apiv1.LightweightRestoreStatus_DATA_RESTORING
	expectedMessage := "kernel restore packed task done"
	expectedProgress := uint32(100)
	expectedBytes := uint64(1024)

	if result.Status != expectedStatus {
		t.Errorf("expected status %v, got %v", expectedStatus, result.Status)
	}

	if result.Message != expectedMessage {
		t.Errorf("expected message %q, got %q", expectedMessage, result.Message)
	}

	if result.DataRestoreProgressPercentage != expectedProgress {
		t.Errorf("expected progress %v, got %v", expectedProgress, result.DataRestoreProgressPercentage)
	}

	if result.DataRestoreBytes != expectedBytes {
		t.Errorf("expected bytes %v, got %v", expectedBytes, result.DataRestoreBytes)
	}
}