package logredactionpolicy

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"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	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 TestUpdateLogRedactionPolicyViaAgent(t *testing.T) {
	cluster, region, tidbgroup := testAgentEntities()
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)

	testCases := []struct {
		name      string
		redactLog bool
		prepare   func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient)
		wantErr   bool
	}{
		{
			name:      "success",
			redactLog: true,
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				setupAgentCredsExpectations(t, regionalCli, encryptor, cluster, region, tidbgroup)
				regionalCli.EXPECT().UpdateLogRedactionPolicyViaAgent(
					gomock.Any(), region, clusterIDStr, "sp-1",
					gomock.AssignableToTypeOf(&regionalserver.UpdateLogRedactionPolicyViaAgentRequest{}),
				).DoAndReturn(func(_ context.Context, _ *entity.Region, _, _ string, req *regionalserver.UpdateLogRedactionPolicyViaAgentRequest) bizerr.BizError {
					require.NotEmpty(t, req.EncryptedUsername)
					require.NotEmpty(t, req.EncryptedPassword)
					require.True(t, req.RedactLog)
					return nil
				})
			},
			wantErr: false,
		},
		{
			name:      "credentials build failure",
			redactLog: false,
			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:      "agent call failure",
			redactLog: true,
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				setupAgentCredsExpectations(t, regionalCli, encryptor, cluster, region, tidbgroup)
				regionalCli.EXPECT().UpdateLogRedactionPolicyViaAgent(
					gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any(),
				).Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "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)

			svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}
			err := svc.UpdateLogRedactionPolicyViaAgent(context.Background(), cluster, region, tidbgroup, tc.redactLog)
			if tc.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestGetLogRedactionPolicyViaAgent(t *testing.T) {
	cluster, region, tidbgroup := testAgentEntities()
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)

	testCases := []struct {
		name      string
		prepare   func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient)
		wantErr   bool
		wantValue bool
	}{
		{
			name: "success redact_log true",
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				setupAgentCredsExpectations(t, regionalCli, encryptor, cluster, region, tidbgroup)
				regionalCli.EXPECT().GetLogRedactionPolicyViaAgent(
					gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any(),
				).Return(&regionalserver.GetLogRedactionPolicyViaAgentResponse{RedactLog: true}, nil)
			},
			wantErr:   false,
			wantValue: true,
		},
		{
			name: "success redact_log false",
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				setupAgentCredsExpectations(t, regionalCli, encryptor, cluster, region, tidbgroup)
				regionalCli.EXPECT().GetLogRedactionPolicyViaAgent(
					gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any(),
				).Return(&regionalserver.GetLogRedactionPolicyViaAgentResponse{RedactLog: false}, nil)
			},
			wantErr:   false,
			wantValue: false,
		},
		{
			name: "credentials build failure",
			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: "agent call failure",
			prepare: func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient) {
				setupAgentCredsExpectations(t, regionalCli, encryptor, cluster, region, tidbgroup)
				regionalCli.EXPECT().GetLogRedactionPolicyViaAgent(
					gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any(),
				).Return(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "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)

			svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}
			got, err := svc.GetLogRedactionPolicyViaAgent(context.Background(), cluster, region, tidbgroup)
			if tc.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, cluster.ClusterId, got.ClusterId)
			require.Equal(t, tc.wantValue, got.GetRedactLog())
		})
	}
}

func setupAgentCredsExpectations(
	t *testing.T,
	regionalCli *mockregionalserver.MockClient,
	encryptor *mockremoteencryptor.MockClient,
	cluster *entity.Cluster,
	region *entity.Region,
	tidbgroup *entity.TidbGroup,
) {
	t.Helper()
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)
	tidbGroupIDStr := strconv.FormatUint(tidbgroup.Id, 10)

	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(mustGenerateTestAgentPublicKeyPEM(t), nil)
	encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).Return("pwd", nil)
}

func testAgentEntities() (*entity.Cluster, *entity.Region, *entity.TidbGroup) {
	return &entity.Cluster{Cluster: &globalngv1.Cluster{ClusterId: 100, EncryptedPlatformAdminPassword: "encrypted"}},
		&entity.Region{Region: &globalcommonv1.Region{State: globalcommonv1.Region_ACTIVE}},
		&entity.TidbGroup{Id: 200}
}

func mustGenerateTestAgentPublicKeyPEM(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})
}