package common

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"strconv"
	"testing"

	"github.com/stretchr/testify/require"
	clusterngv1beta "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"go.uber.org/mock/gomock"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	mockregionalserver "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver/mock_regionalserver"
	mockremoteencryptor "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/remoteencryptor/mock_remoteencryptor"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

func TestBuildClusterAgentCredentials(t *testing.T) {
	cluster := &entity.Cluster{Cluster: &globalngv1.Cluster{ClusterId: 100, EncryptedPlatformAdminPassword: "encrypted"}}
	region := &entity.Region{Region: &globalcommonv1.Region{State: globalcommonv1.Region_ACTIVE}}
	tidbgroup := &entity.TidbGroup{Id: 200}
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)
	tidbGroupIDStr := strconv.FormatUint(tidbgroup.Id, 10)

	testCases := []struct {
		name    string
		prepare func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient)
		wantErr bool
	}{
		{
			name: "success",
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, &clusterngv1beta.GetTiDBGroupRequest{
					ClusterId:   clusterIDStr,
					TidbGroupId: tidbGroupIDStr,
				}).Return(&clusterngv1beta.TiDBGroup{Id: "default", SharedPoolId: "sp-1"}, nil)
				regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
					Return(mustGenerateTestPublicKeyPEM(t), nil)
				encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).Return("pwd", nil)
			},
			wantErr: false,
		},
		{
			name: "get tidb group failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, _ *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).
					Return(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))
			},
			wantErr: true,
		},
		{
			name: "get public key failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, _ *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).
					Return(&clusterngv1beta.TiDBGroup{Id: "default", SharedPoolId: "sp-1"}, nil)
				regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
					Return(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))
			},
			wantErr: true,
		},
		{
			name: "parse public key failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, _ *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).
					Return(&clusterngv1beta.TiDBGroup{Id: "default", SharedPoolId: "sp-1"}, nil)
				regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
					Return([]byte("not-a-pem"), nil)
			},
			wantErr: true,
		},
		{
			name: "decrypt password failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).
					Return(&clusterngv1beta.TiDBGroup{Id: "default", SharedPoolId: "sp-1"}, nil)
				regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
					Return(mustGenerateTestPublicKeyPEM(t), nil)
				encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).
					Return("", bizerr.New(bizerr.ErrCodeUnknown, "failed"))
			},
			wantErr: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			regionalCli := mockregionalserver.NewMockClient(ctrl)
			encryptor := mockremoteencryptor.NewMockClient(ctrl)
			tc.prepare(regionalCli, encryptor)

			cred, err := BuildClusterAgentCredentials(context.Background(), regionalCli, encryptor, cluster, region, tidbgroup)
			if tc.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, "sp-1", cred.SharedPoolID)
			require.NotEmpty(t, cred.EncryptedUsername)
			require.NotEmpty(t, cred.EncryptedPassword)
		})
	}
}

func mustGenerateTestPublicKeyPEM(t *testing.T) []byte {
	t.Helper()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	pubBytes := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)
	return pem.EncodeToMemory(&pem.Block{Type: "RSA PUBLIC KEY", Bytes: pubBytes})
}