package auditlog

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"
	"google.golang.org/protobuf/proto"

	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"
	svccommon "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/common"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

func TestUpdateAuditLogSettingViaAgent_Success(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	regionalCli := mockregionalserver.NewMockClient(ctrl)
	encryptor := mockremoteencryptor.NewMockClient(ctrl)
	svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}

	cluster, region, tidbgroup := testAuditAgentEntities()
	publicKeyPEM := mustGeneratePublicKeyPEM(t)
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)

	regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, &clusterngv1beta.GetTiDBGroupRequest{
		ClusterId:   clusterIDStr,
		TidbGroupId: strconv.FormatUint(tidbgroup.Id, 10),
	}).Return(&clusterngv1beta.TiDBGroup{Id: "default", SharedPoolId: "sp-1"}, nil)
	regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
		Return(publicKeyPEM, nil)
	encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).Return("pwd", nil)
	regionalCli.EXPECT().UpdateAuditLogSettingViaAgent(
		gomock.Any(),
		region,
		clusterIDStr,
		"sp-1",
		gomock.AssignableToTypeOf(&regionalserver.UpdateAuditLogSettingViaAgentRequest{}),
	).DoAndReturn(func(_ context.Context, _ *entity.Region, _, _ string, req *regionalserver.UpdateAuditLogSettingViaAgentRequest) bizerr.BizError {
		require.NotEmpty(t, req.EncryptedUsername)
		require.NotEmpty(t, req.EncryptedPassword)
		require.True(t, req.Enabled)
		require.Equal(t, int32(64), req.MaxSizeMib)
		return nil
	})

	err := svc.UpdateAuditLogSettingViaAgent(context.Background(), cluster, region, tidbgroup, &globalngv1.AuditLogSetting{
		Enabled:      proto.Bool(true),
		Unredacted:   proto.Bool(true),
		Format:       "JSON",
		MaxSizeMib:   64,
		MaxAgeSecond: proto.Int32(3600),
	})
	require.NoError(t, err)
}

func TestBuildClusterAgentAuditCredentials_Failures(t *testing.T) {
	cluster, region, tidbgroup := testAuditAgentEntities()
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)
	tidbGroupIDStr := strconv.FormatUint(tidbgroup.Id, 10)

	testCases := []struct {
		name    string
		prepare func(regionalCli *mockregionalserver.MockClient, encryptor *mockremoteencryptor.MockClient)
	}{
		{
			name: "get tidb group failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, _ *mockremoteencryptor.MockClient) {
				regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, &clusterngv1beta.GetTiDBGroupRequest{
					ClusterId:   clusterIDStr,
					TidbGroupId: tidbGroupIDStr,
				}).Return(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))
			},
		},
		{
			name: "get public key failed",
			prepare: func(regionalCli *mockregionalserver.MockClient, _ *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(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))
			},
		},
		{
			name: "decrypt password failed",
			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(mustGeneratePublicKeyPEM(t), nil)
				encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).
					Return("", bizerr.New(bizerr.ErrCodeUnknown, "failed"))
			},
		},
	}

	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)
			_, err := svccommon.BuildClusterAgentCredentials(context.Background(), regionalCli, encryptor, cluster, region, tidbgroup)
			require.Error(t, err)
		})
	}
}

func TestCreateDeleteUpdateAuditLogFilterRuleViaAgent_AgentError(t *testing.T) {
	cluster, region, tidbgroup := testAuditAgentEntities()
	clusterIDStr := strconv.FormatUint(cluster.ClusterId, 10)
	tidbGroupIDStr := strconv.FormatUint(tidbgroup.Id, 10)
	publicKeyPEM := mustGeneratePublicKeyPEM(t)

	setupCommon := 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).AnyTimes()
		regionalCli.EXPECT().GetClusterAgentPublicKey(gomock.Any(), region, clusterIDStr, "sp-1").
			Return(publicKeyPEM, nil).AnyTimes()
		encryptor.EXPECT().DecryptBase64(gomock.Any(), cluster.EncryptedPlatformAdminPassword).Return("pwd", nil).AnyTimes()
	}

	t.Run("create failed", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		defer ctrl.Finish()

		regionalCli := mockregionalserver.NewMockClient(ctrl)
		encryptor := mockremoteencryptor.NewMockClient(ctrl)
		svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}
		setupCommon(regionalCli, encryptor)
		regionalCli.EXPECT().CreateAuditLogFilterRuleViaAgent(gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))

		err := svc.CreateAuditLogFilterRuleViaAgent(context.Background(), cluster, region, tidbgroup, &globalngv1.AuditLogFilterRule{UserExpr: "u1"})
		require.Error(t, err)
	})

	t.Run("delete failed", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		defer ctrl.Finish()

		regionalCli := mockregionalserver.NewMockClient(ctrl)
		encryptor := mockremoteencryptor.NewMockClient(ctrl)
		svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}
		setupCommon(regionalCli, encryptor)
		regionalCli.EXPECT().DeleteAuditLogFilterRuleViaAgent(gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))

		err := svc.DeleteAuditLogFilterRuleViaAgent(context.Background(), cluster, region, tidbgroup, "u1")
		require.Error(t, err)
	})

	t.Run("update failed", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		defer ctrl.Finish()

		regionalCli := mockregionalserver.NewMockClient(ctrl)
		encryptor := mockremoteencryptor.NewMockClient(ctrl)
		svc := &service{regionalServerCli: regionalCli, encryptor: encryptor}
		setupCommon(regionalCli, encryptor)
		regionalCli.EXPECT().UpdateAuditLogFilterRuleViaAgent(gomock.Any(), region, clusterIDStr, "sp-1", gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "failed"))

		err := svc.UpdateAuditLogFilterRuleViaAgent(context.Background(), cluster, region, tidbgroup, &globalngv1.AuditLogFilterRule{UserExpr: "u1"})
		require.Error(t, err)
	})
}

func testAuditAgentEntities() (*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 mustGeneratePublicKeyPEM(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})
}