package server

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func TestGetImportRoleConfig(t *testing.T) {
	tests := []struct {
		name          string
		clusterID     string
		setupMocks    func(*repomocks.TiDBClusterRepo, *repomocks.SharedPoolRepo, *dataplanemgrmocks.Client)
		expectedCode  codes.Code
		expectedError string
		validate      func(*testing.T, *entity.ImportRoleConfig, error)
	}{
		{
			name:      "successfully get import role config",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetKeyspaceRoleRequest) bool {
					return req.GetDataplaneId() == "dp-1" &&
						req.GetKeyspaceName() == "123" &&
						req.GetBucketId() == "sp-1-us-west-2-internal" &&
						req.GetAllowedPrefix() == "internal/prefix/123"
				})).Return(readyImportKeyspaceRole(), nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "arn:aws:iam::123456789012:role/import-role-123", result.GetRoleArn())
				assert.Equal(t, "123", result.GetExternalId())
				assert.Equal(t, "physical-internal-bucket", result.GetBucket())
				assert.Equal(t, "internal/prefix/123", result.GetAllowedPrefix())
				assert.Equal(t, "s3://physical-internal-bucket/internal/prefix/123", result.GetCloudStorageUri())
			},
		},
		{
			name:      "normalizes internal prefix for returned cloud storage uri",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.InternalStorage.Prefix = " /internal/prefix/ "
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(readyImportKeyspaceRole(), nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "internal/prefix/123", result.GetAllowedPrefix())
				assert.Equal(t, "s3://physical-internal-bucket/internal/prefix/123", result.GetCloudStorageUri())
			},
		},
		{
			name:      "builds cloud storage uri without internal prefix",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.InternalStorage.Prefix = ""
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
				setupActiveDataplane(dataplaneClient)
				role := readyImportKeyspaceRole()
				role.AllowedPrefix = "123"
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(role, nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "123", result.GetAllowedPrefix())
				assert.Equal(t, "s3://physical-internal-bucket/123", result.GetCloudStorageUri())
			},
		},
		{
			name:      "builds alicloud cloud storage uri with normalized prefix",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.Provider.Type = apiv1.ProviderType_ALICLOUD
				sharedPool.Provider.Region = "ap-southeast-1"
				sharedPool.InternalStorage.Prefix = " /internal/prefix/ "
				sharedPool.InternalStorage.Region = "ap-southeast-1"
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
				setupActiveDataplane(dataplaneClient)
				role := readyImportKeyspaceRole()
				role.BucketId = "sp-1-ap-southeast-1-internal"
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(role, nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "internal/prefix/123", result.GetAllowedPrefix())
				assert.Equal(t, "oss://physical-internal-bucket/internal/prefix/123", result.GetCloudStorageUri())
			},
		},
		{
			name:          "empty cluster id",
			clusterID:     "",
			setupMocks:    func(*repomocks.TiDBClusterRepo, *repomocks.SharedPoolRepo, *dataplanemgrmocks.Client) {},
			expectedCode:  codes.InvalidArgument,
			expectedError: "tidbCluster id is required",
		},
		{
			name:      "cluster not found",
			clusterID: "404",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				tidbClusterRepo.On("Load", mock.Anything, "404", false).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedCode:  codes.NotFound,
			expectedError: "tidbCluster 404 not found",
		},
		{
			name:      "cluster load error",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				tidbClusterRepo.On("Load", mock.Anything, "123", false).Return(nil, errors.New("database error"))
			},
			expectedCode:  codes.Internal,
			expectedError: "failed to load tidbCluster: database error",
		},
		{
			name:      "shared pool not found",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				tidbClusterRepo.On("Load", mock.Anything, "123", false).Return(&entity.TiDBCluster{
					Id:           "123",
					SharedPoolId: "404",
				}, nil)
				sharedPoolRepo.On("Load", mock.Anything, "404", false).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedCode:  codes.NotFound,
			expectedError: "sharedPool 404 not found",
		},
		{
			name:      "topology driven shared pool with user keyspace is supported",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.CapacityMode = apiv1.CapacityMode_TOPOLOGY_DRIVEN
				sharedPool.ComponentGroups = []*apiv1.Component{{
					ComponentType: apiv1.ComponentType_TIDB_WORKER,
					NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
				}}
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(readyImportKeyspaceRole(), nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "internal/prefix/123", result.GetAllowedPrefix())
			},
		},
		{
			name:      "topology driven shared pool without TiDB worker uses cluster keyspace",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.CapacityMode = apiv1.CapacityMode_TOPOLOGY_DRIVEN
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
				setupActiveDataplane(dataplaneClient)
				role := readyImportKeyspaceRole()
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetKeyspaceRoleRequest) bool {
					return req.GetDataplaneId() == "dp-1" &&
						req.GetKeyspaceName() == "123" &&
						req.GetBucketId() == "sp-1-us-west-2-internal" &&
						req.GetAllowedPrefix() == "internal/prefix/123" &&
						req.GetExternalId() == "123"
				})).Return(role, nil)
			},
			validate: func(t *testing.T, result *entity.ImportRoleConfig, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "123", result.GetExternalId())
				assert.Equal(t, "internal/prefix/123", result.GetAllowedPrefix())
				assert.Equal(t, "s3://physical-internal-bucket/internal/prefix/123", result.GetCloudStorageUri())
			},
		},
		{
			name:      "unsupported provider",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.Provider.Type = apiv1.ProviderType_GCP
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "not supported for provider GCP",
		},
		{
			name:      "internal storage is not configured",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.InternalStorage = nil
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "internal storage is not configured",
		},
		{
			name:      "internal bucket is not ready",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.InternalStorage.Bucket = ""
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "internal bucket is not ready",
		},
		{
			name:      "internal storage region is not configured",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				sharedPool := baseImportRoleSharedPool()
				sharedPool.InternalStorage.Region = ""
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, sharedPool)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "internal storage region is not configured",
		},
		{
			name:      "dataplane is not active",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				dataplaneClient.On("GetDataPlane", mock.Anything, mock.Anything).Return(&dplanemgrv1.DataPlane{
					Status: dplanemgrv1.DataPlane_Creating,
				}, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "dataplane dp-1 is not active: Creating",
		},
		{
			name:      "keyspace role not found",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "role not found",
		},
		{
			name:      "keyspace role is creating",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(&dplanemgrv1.KeyspaceRole{
					Status:  dplanemgrv1.KeyspaceRole_Creating,
					Message: "keyspace role ARN is not ready",
				}, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "keyspace role ARN is not ready",
		},
		{
			name:      "keyspace role is deleting",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(&dplanemgrv1.KeyspaceRole{
					Status: dplanemgrv1.KeyspaceRole_Deleting,
				}, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "Deleting",
		},
		{
			name:      "keyspace role failed",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(&dplanemgrv1.KeyspaceRole{
					Status:  dplanemgrv1.KeyspaceRole_Failed,
					Message: "policy reconcile failed",
				}, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "policy reconcile failed",
		},
		{
			name:      "ready keyspace role has empty ARN",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(&dplanemgrv1.KeyspaceRole{
					Status: dplanemgrv1.KeyspaceRole_Ready,
				}, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "role ARN is empty",
		},
		{
			name:      "ready keyspace role contract does not match derived config",
			clusterID: "123",
			setupMocks: func(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneClient *dataplanemgrmocks.Client) {
				setupImportRoleConfigRepos(tidbClusterRepo, sharedPoolRepo, baseImportRoleSharedPool())
				setupActiveDataplane(dataplaneClient)
				role := readyImportKeyspaceRole()
				role.AllowedPrefix = "other-prefix/123"
				dataplaneClient.On("GetKeyspaceRole", mock.Anything, mock.Anything).Return(role, nil)
			},
			expectedCode:  codes.FailedPrecondition,
			expectedError: "allowed_prefix",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
			sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
			dataplaneClient := dataplanemgrmocks.NewClient(t)
			tt.setupMocks(tidbClusterRepo, sharedPoolRepo, dataplaneClient)

			server := &serverImpl{
				tidbClusterRepo: tidbClusterRepo,
				sharedPoolRepo:  sharedPoolRepo,
				dataplaneClient: dataplaneClient,
			}

			result, err := server.GetImportRoleConfig(context.Background(), tt.clusterID)
			if tt.expectedCode != codes.OK {
				require.Error(t, err)
				assert.Equal(t, tt.expectedCode, status.Code(err))
				assert.Contains(t, err.Error(), tt.expectedError)
				assert.Nil(t, result)
				return
			}
			tt.validate(t, result, err)
		})
	}
}

func setupImportRoleConfigRepos(tidbClusterRepo *repomocks.TiDBClusterRepo, sharedPoolRepo *repomocks.SharedPoolRepo, sharedPool *entity.SharedPool) {
	tidbClusterRepo.On("Load", mock.Anything, "123", false).Return(&entity.TiDBCluster{
		Id:           "123",
		SharedPoolId: sharedPool.GetId(),
	}, nil)
	sharedPoolRepo.On("Load", mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil)
}

func setupActiveDataplane(dataplaneClient *dataplanemgrmocks.Client) {
	dataplaneClient.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dp-1"
	})).Return(&dplanemgrv1.DataPlane{
		Status: dplanemgrv1.DataPlane_Active,
	}, nil)
}

func readyImportKeyspaceRole() *dplanemgrv1.KeyspaceRole {
	return &dplanemgrv1.KeyspaceRole{
		RoleArn:       "arn:aws:iam::123456789012:role/import-role-123",
		ExternalId:    "123",
		BucketId:      "sp-1-us-west-2-internal",
		AllowedPrefix: "internal/prefix/123",
		Status:        dplanemgrv1.KeyspaceRole_Ready,
	}
}

func baseImportRoleSharedPool() *entity.SharedPool {
	return &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "sp-1",
			DataplaneName: "dp-1",
			CapacityMode:  apiv1.CapacityMode_RCU_DRIVEN,
			Provider: &apiv1.Provider{
				Type:   apiv1.ProviderType_AWS,
				Region: "us-west-2",
			},
			InternalStorage: &apiv1.SharedPool_ObjectStorage{
				Bucket: "physical-internal-bucket",
				Prefix: "internal/prefix",
				Region: "us-west-2",
			},
		},
	}
}