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"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor"
	executormocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func TestSetAuditLogSetting(t *testing.T) {
	validAWSRoleARN := "arn:aws:iam::123456789012:role/test-role"
	validAliCloudRoleARN := "acs:ram::123456789012:role/test-role"

	setupSuccessfulMocks := func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
		// Mock tidb cluster load
		tidbCluster := &entity.TiDBCluster{
			Id:    "testclusterid",
			Phase: apiv1.Cluster_AVAILABLE,
		}
		tidbClusterRepo.On("Load", mock.Anything, "testclusterid", false).Return(tidbCluster, nil)

		// Mock tidb groups list
		tidbGroups := []*entity.TiDBGroup{
			{
				TiDBGroup: apiv1.TiDBGroup{
					Id:           "test-group-id",
					SharedPoolId: "testsharedpool-id",
				},
			},
		}
		tidbGroupRepo.On("List", mock.Anything, mock.Anything, mock.Anything, []string{"testclusterid"}, mock.Anything, false).Return(tidbGroups, nil)

		// Mock shared pool load
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id: "testsharedpool-id",
				ComponentGroups: []*apiv1.Component{
					{
						Id:            "test-group-id",
						ComponentType: apiv1.ComponentType_TIDB,
						Phase:         apiv1.Component_AVAILABLE,
					},
				},
			},
		}
		sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-id", false).Return(sharedPool, nil)

		// Mock merge status (no running DAGs)
		dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

		// Mock lock/unlock operations
		lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)

		// Mock no running tidb group DAGs
		dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

		// Mock tidb cluster save
		tidbClusterRepo.On("Save", mock.Anything, mock.Anything).Return(nil)

		// Mock DAG submit
		dagExecutor.On("Submit", mock.Anything, dag.PlanUpdateTiDBCluster, mock.Anything, mock.Anything).Return(nil)
	}

	tests := []struct {
		name            string
		tidbClusterId   string
		auditLogSetting *entity.AuditLogSetting
		cloudProvider   string
		setupMocks      func(*repomnocks.TiDBClusterRepo, *repomnocks.TiDBGroupRepo, *repomnocks.SharedPoolRepo, *repomnocks.LockRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager)
		expectedError   string
		expectedCode    codes.Code
		validate        func(*testing.T, *entity.AuditLogSetting, error)
	}{
		{
			name:          "successful set audit log setting for aws",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks:    setupSuccessfulMocks,
			validate: func(t *testing.T, result *entity.AuditLogSetting, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, validAWSRoleARN, result.GetProviderIdentity().GetCloudRoleIdentity())
			},
		},
		{
			name:          "successful set audit log setting for alicloud",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAliCloudRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_ALICLOUD.String(),
			setupMocks:    setupSuccessfulMocks,
			validate: func(t *testing.T, result *entity.AuditLogSetting, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, validAliCloudRoleARN, result.GetProviderIdentity().GetCloudRoleIdentity())
			},
		},
		{
			name:          "empty tidb cluster id",
			tidbClusterId: "",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "tidb cluster id is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:            "nil audit log setting",
			tidbClusterId:   "testclusterid",
			auditLogSetting: nil,
			cloudProvider:   apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "audit log setting is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "invalid AWS role ARN",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: "invalid-arn",
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "invalid aws role arn",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "invalid alicloud role arn",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: "invalid-arn",
				},
			},
			cloudProvider: apiv1.ProviderType_ALICLOUD.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "invalid alicloud role arn",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "GCP provider not supported",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: "gcp-identity",
				},
			},
			cloudProvider: apiv1.ProviderType_GCP.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "GCP cloud provider is not supported yet",
			expectedCode:  codes.Unimplemented,
		},
		{
			name:          "tidb cluster not found",
			tidbClusterId: "nonexistentid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				tidbClusterRepo.On("Load", mock.Anything, "nonexistentid", false).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedError: "tidb cluster nonexistentid not found",
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:          "tidb cluster is creating",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				tidbCluster := &entity.TiDBCluster{
					Id:    "testclusterid",
					Phase: apiv1.Cluster_CREATING,
				}
				tidbClusterRepo.On("Load", mock.Anything, "testclusterid", false).Return(tidbCluster, nil)
			},
			expectedError: "tidb cluster testclusterid is creating, please try again later",
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:          "no associated tidb groups",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				tidbCluster := &entity.TiDBCluster{
					Id:    "testclusterid",
					Phase: apiv1.Cluster_AVAILABLE,
				}
				tidbClusterRepo.On("Load", mock.Anything, "testclusterid", false).Return(tidbCluster, nil)

				// Mock empty tidb groups list
				tidbGroupRepo.On("List", mock.Anything, mock.Anything, mock.Anything, []string{"testclusterid"}, mock.Anything, false).Return([]*entity.TiDBGroup{}, nil)
			},
			expectedError: "tidb cluster testclusterid has no associated tidb group",
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:          "tidb group is creating",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				tidbCluster := &entity.TiDBCluster{
					Id:    "testclusterid",
					Phase: apiv1.Cluster_AVAILABLE,
				}
				tidbClusterRepo.On("Load", mock.Anything, "testclusterid", false).Return(tidbCluster, nil)

				tidbGroups := []*entity.TiDBGroup{
					{
						TiDBGroup: apiv1.TiDBGroup{
							Id:           "test-group-id",
							SharedPoolId: "testsharedpool-id",
						},
					},
				}
				tidbGroupRepo.On("List", mock.Anything, mock.Anything, mock.Anything, []string{"testclusterid"}, mock.Anything, false).Return(tidbGroups, nil)

				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpool-id",
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-id", false).Return(sharedPool, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)

				// Mock running create tidb group DAG for the specific tidb group
				runningDags := executor.NewExecTaskIdStatus()
				runningDags["test-dag-id"] = &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							RootName: dag.PlanCreateTiDBGroup,
							Args: actions.ActionArgs{
								workflow.SharedPoolIdKey: actions.ActionArgString("testsharedpool-id"),
								workflow.ComponentIdsKey: actions.ActionArgString("test-group-id"),
							},
						},
					},
					State: executor.ExecTaskStateTypeRunningOK,
				}

				// Mock specific DAG query for TiDB group creation check with running DAGs
				dagExecutor.On("QueryRunnings", mock.Anything, "create-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(runningDags, nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "update-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "delete-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)

				// Mock merge status for shared pool (no running DAGs for merge) - this should come after the specific mocks
				emptyDags := executor.NewExecTaskIdStatus()
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(emptyDags, nil).Maybe()
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe()
			},
			expectedError: "tidb group test-group-id is creating, please try again later",
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:          "tidb cluster save error",
			tidbClusterId: "testclusterid",
			auditLogSetting: &entity.AuditLogSetting{
				ProviderIdentity: &entity.ProviderIdentity{
					CloudRoleIdentity: validAWSRoleARN,
				},
			},
			cloudProvider: apiv1.ProviderType_AWS.String(),
			setupMocks: func(tidbClusterRepo *repomnocks.TiDBClusterRepo, tidbGroupRepo *repomnocks.TiDBGroupRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				tidbCluster := &entity.TiDBCluster{
					Id:    "testclusterid",
					Phase: apiv1.Cluster_AVAILABLE,
				}
				tidbClusterRepo.On("Load", mock.Anything, "testclusterid", false).Return(tidbCluster, nil)

				tidbGroups := []*entity.TiDBGroup{
					{
						TiDBGroup: apiv1.TiDBGroup{
							Id:           "test-group-id",
							SharedPoolId: "testsharedpool-id",
						},
					},
				}
				tidbGroupRepo.On("List", mock.Anything, mock.Anything, mock.Anything, []string{"testclusterid"}, mock.Anything, false).Return(tidbGroups, nil)

				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpool-id",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "test-group-id",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-id", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe()

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)

				// Mock no running tidb group DAGs
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)

				// Mock tidb cluster save error
				tidbClusterRepo.On("Save", mock.Anything, mock.Anything).Return(errors.New("save error"))
			},
			expectedError: "failed to save tidb cluster: save error",
			expectedCode:  codes.Internal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
			tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t)
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			lockRepo := repomnocks.NewLockRepo(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scaleManager := bractionmocks.NewScaleManager(t)
			// Setup mocks
			tt.setupMocks(tidbClusterRepo, tidbGroupRepo, sharedPoolRepo, lockRepo, dagExecutor, scaleManager)

			// Create server implementation
			server := &serverImpl{
				cloudProvider:   tt.cloudProvider,
				tidbClusterRepo: tidbClusterRepo,
				tidbGroupRepo:   tidbGroupRepo,
				sharedPoolRepo:  sharedPoolRepo,
				lockRepo:        lockRepo,
				dagExecutor:     dagExecutor,
				transactioner:   &mockTransactioner{},
				scaleManager:    scaleManager,
			}

			// Execute test
			result, err := server.SetAuditLogSetting(context.Background(), tt.tidbClusterId, tt.auditLogSetting)

			switch {
			case tt.expectedError != "":
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.expectedError)
				if tt.expectedCode != codes.OK {
					assert.Equal(t, tt.expectedCode, status.Code(err))
				}
				assert.Nil(t, result)

			case tt.validate != nil:
				tt.validate(t, result, err)

			default:
				require.NoError(t, err)
				require.NotNil(t, result)
			}
		})
	}
}

func TestValidateAliCloudARN(t *testing.T) {
	tests := []struct {
		name     string
		arn      string
		expected bool
	}{
		{
			name:     "valid ram role without region",
			arn:      "acs:ram::123456789012:role/test-role",
			expected: true,
		},
		{
			name:     "valid ecs instance with region",
			arn:      "acs:ecs:cn-hangzhou:987654321098:instance/ecs-123456",
			expected: true,
		},
		{
			name:     "empty arn",
			arn:      "",
			expected: false,
		},
		{
			name:     "non acs prefix",
			arn:      "arn:aws:iam::123456789012:role/test-role",
			expected: false,
		},
		{
			name:     "non numeric account",
			arn:      "acs:ram::abc123:role/test-role",
			expected: false,
		},
		{
			name:     "missing resource",
			arn:      "acs:ram::123456789012:",
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := utils.ValidateAliCloudARN(tt.arn)
			assert.Equal(t, tt.expected, result)
		})
	}
}