package auditlogv2

import (
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"cloud.google.com/go/storage"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/samber/lo"
	"github.com/spf13/viper"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"

	openapicommonv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/common/v1beta1"
	dedicatedapi "github.com/tidbcloud/tidb-management-service/functional-test/pkg/dedicatedapi/api"
	"github.com/tidbcloud/tidb-management-service/functional-test/pkg/dedicatedapi/api/client"
	"github.com/tidbcloud/tidb-management-service/functional-test/pkg/sql"
	"github.com/tidbcloud/tidb-management-service/utilities/randstr"
)

// TestAuditLogV2CustomerBucket verifies end-to-end behavior of tidb cluster setting storing audit log to customer's object Storage.
func TestAuditLogV2CustomerBucket(t *testing.T) {
	tester := &auditLogV2CustomerBucketTest{
		T:                    t,
		host:                 viper.GetString("host"),
		apiKey:               viper.GetString("api_key"),
		prjID:                viper.GetString("prj_id"),
		provider:             viper.GetString("cloud_provider"),
		regionCode:           viper.GetString("region_code"),
		clusterId:            viper.GetString("cluster_id"),
		clusterNamePrefix:    viper.GetString("cluster_name_prefix"),
		awsS3BucketName:      viper.GetString("aws_s3_bucket_name"),
		awsS3RoleArn:         viper.GetString("aws_s3_role_arn"),
		awsProfile:           viper.GetString("aws_profile"),
		azureBlobAccountName: viper.GetString("azure_blob_account_name"),
		azureBlobContainer:   viper.GetString("azure_blob_container"),
		azureBlobSASToken:    viper.GetString("azure_blob_sas_token"),
		gcpGCSBucketName:     viper.GetString("gcp_gcs_bucket_name"),
		gcpCredentialsFile:   viper.GetString("gcp_credentials_file"),
		port:                 4000,
		password:             randstr.GenWithCharset(16, randstr.AlphanumericCharset),
		tidbNodeSpec:         &client.V1beta1NodeSpec{NodeSpecKey: lo.ToPtr("2C8G")},
		tikvNodeSpec: &client.V1beta1NodeSpec{
			NodeSpecKey:          lo.ToPtr("2C8G"),
			DefaultStorageSizeGi: lo.ToPtr(int32(200)),
		},
		filterRuleIds: make([]string, 0),
	}

	tester.useExistingCluster = tester.clusterId != ""
	tester.clusterName = tester.clusterNamePrefix + randstr.GenWithCharset(4, randstr.AlphabetCharset)
	tester.cli = dedicatedapi.NewClientT(tester.T, tester.host, tester.apiKey)

	tests := []struct {
		name string
		fn   func(t *testing.T)
	}{
		{"00_EnsureNetworkContainer", tester.Test_00_EnsureNetworkContainer},
		{"01_CreateCluster", tester.Test_01_CreateCluster},
		{"02_PollClusterUntilActive", tester.Test_02_PollClusterUntilActive},
		{"03_ResetRootPassword", tester.Test_03_ResetRootPassword},
		{"04_EnableAuditLogToCustomer", tester.Test_04_EnableAuditLogToCustomer},
		{"05_SetLogFormatToJSON", tester.Test_05_SetLogFormatToJSON},
		{"06_ConfigureLogRotationBySize", tester.Test_06_ConfigureLogRotationBySize},
		{"07_ConfigureLogRotationByTime", tester.Test_07_ConfigureLogRotationByTime},
		{"08_DisableLogRedaction", tester.Test_08_DisableLogRedaction},
		{"09_CreateFilterRules", tester.Test_09_CreateFilterRules},
		{"10_UpdateFilterRule", tester.Test_10_UpdateFilterRule},
		{"11_DisableFilterRule", tester.Test_11_DisableFilterRule},
		{"12_EnableFilterRule", tester.Test_12_EnableFilterRule},
		{"13_DeleteFilterRule", tester.Test_13_DeleteFilterRule},
		{"14_SetPublicNetworkAccess", tester.Test_14_SetPublicNetworkAccess},
		{"15_ExecuteSql", tester.Test_15_ExecuteSql},
		{"16_PauseCluster", tester.Test_16_PauseCluster},
		{"17_CheckAuditLogsInCustomerBucket", tester.Test_17_CheckAuditLogsInCustomerBucket},
		{"18_DisableAuditLog", tester.Test_18_DisableAuditLog},
	}

	for _, tt := range tests {
		if t.Failed() {
			t.Logf("Skipping test %s due to previous failure.", tt.name)
			continue
		}
		t.Run(tt.name, tt.fn)
	}
}

type auditLogV2CustomerBucketTest struct {
	*testing.T

	host                 string
	apiKey               string
	prjID                string
	clusterNamePrefix    string
	provider             string
	regionCode           string
	password             string
	clusterName          string
	port                 int32
	awsS3BucketName      string
	awsS3RoleArn         string
	awsProfile           string
	azureBlobAccountName string
	azureBlobContainer   string
	azureBlobSASToken    string
	gcpGCSBucketName     string
	gcpCredentialsFile   string
	cli                  *dedicatedapi.Cli
	tidbNodeSpec         *client.V1beta1NodeSpec
	tikvNodeSpec         *client.V1beta1NodeSpec
	clusterId            string
	useExistingCluster   bool
	filterRuleIds        []string
}

func (su *auditLogV2CustomerBucketTest) Test_00_EnsureNetworkContainer(t *testing.T) {
	req := su.cli.NetworkContainerServiceAPI.NetworkContainerServiceListNetworkContainers(t.Context()).
		ProjectId(su.prjID).
		CloudProvider(su.provider).
		PageSize(pageSize)
	resp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to list network containers: %v", err)
	}
	isProjectAndCloudProviderValidForAll := lo.Reduce(
		resp.GetNetworkContainers(),
		func(agg bool, item client.V1beta1NetworkContainer, _ int) bool {
			pid, ok := item.GetLabels()["tidb.cloud/project"]
			return agg &&
				ok && pid == su.prjID && // Check result's project ID.
				string(item.GetCloudProvider()) == su.provider // Check result's cloud provider.
		},
		true,
	)
	if !isProjectAndCloudProviderValidForAll {
		t.Fatal("network container check failed")
	}
}

func (su *auditLogV2CustomerBucketTest) Test_01_CreateCluster(t *testing.T) {
	if su.useExistingCluster {
		t.Skip("Skipping cluster creation, using existing cluster.")
	}

	// Create a cluster.
	req := su.cli.ClusterServiceAPI.ClusterServiceCreateCluster(t.Context()).
		Cluster(client.TidbCloudOpenApidedicatedv1beta1Cluster{
			DisplayName: su.clusterName,
			RegionId:    fmt.Sprintf("%s-%s", su.provider, su.regionCode),
			Labels:      lo.ToPtr(map[string]string{"tidb.cloud/project": su.prjID}),
			TidbNodeSetting: client.V1beta1ClusterTidbNodeSetting{
				NodeSpecKey: lo.ToPtr(su.tidbNodeSpec.GetNodeSpecKey()),
				TidbNodeGroups: []client.Dedicatedv1beta1TidbNodeGroup{
					{
						NodeCount: 1,
					},
				},
			},
			TikvNodeSetting: client.V1beta1ClusterStorageNodeSetting{
				NodeCount:     3,
				NodeSpecKey:   su.tikvNodeSpec.GetNodeSpecKey(),
				StorageSizeGi: su.tikvNodeSpec.GetDefaultStorageSizeGi(),
			},
			Port:         su.port,
			RootPassword: lo.ToPtr(su.password),
		})

	resp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to create cluster: %v", err)
	}
	if resp.GetDisplayName() != su.clusterName {
		t.Fatalf("expected display name %s, got %s", su.clusterName, resp.GetDisplayName())
	}
	if resp.GetRegionId() != fmt.Sprintf("%s-%s", su.provider, su.regionCode) {
		t.Fatalf("expected region id %s, got %s", fmt.Sprintf("%s-%s", su.provider, su.regionCode), resp.GetRegionId())
	}
	if resp.GetState() != client.COMMONV1BETA1CLUSTERSTATE_CREATING {
		t.Fatalf("expected cluster state CREATING, got %s", resp.GetState())
	}

	su.clusterId = resp.GetClusterId()
}

func (su *auditLogV2CustomerBucketTest) Test_02_PollClusterUntilActive(t *testing.T) {
	if su.useExistingCluster {
		t.Skip("Skipping poll, assuming existing cluster is active.")
	}

	if su.clusterId == "" {
		t.Fatal("clusterId is not set, cannot poll cluster.")
	}

	// List clusters to find the cluster.
	req := su.cli.ClusterServiceAPI.ClusterServiceListClusters(t.Context()).
		ProjectId(su.prjID).
		PageSize(pageSize)
	resp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to list clusters: %v", err)
	}
	if !lo.ContainsBy(resp.Clusters, func(item client.TidbCloudOpenApidedicatedv1beta1Cluster) bool {
		return item.GetClusterId() == su.clusterId
	}) {
		t.Fatalf("cluster %s not found in list", su.clusterId)
	}

	// Poll the cluster using GET method until it becomes active.
	// Timeout: 30 minutes (1800 seconds).
	for range 180 {
		req := su.cli.ClusterServiceAPI.ClusterServiceGetCluster(t.Context(), su.clusterId)
		resp, _, err := req.Execute()
		if err != nil {
			t.Logf("failed to get cluster: %v", err)
			time.Sleep(10 * time.Second)
			continue
		}
		if resp == nil {
			t.Logf("cluster is not ready yet")
			time.Sleep(10 * time.Second)
			continue
		}

		switch resp.GetState() {
		case client.COMMONV1BETA1CLUSTERSTATE_ACTIVE:
			return
		case client.COMMONV1BETA1CLUSTERSTATE_CREATING:
			t.Logf("Cluster state is CREATING")
		default:
			t.Fatalf("Cluster state is not ACTIVE or CREATING, but %s", resp.GetState())
		}
		time.Sleep(10 * time.Second)
	}

	t.Fatal("Cluster did not become active within the timeout period")
}

func (su *auditLogV2CustomerBucketTest) Test_03_ResetRootPassword(t *testing.T) {
	req := su.cli.ClusterServiceAPI.ClusterServiceResetRootPassword(t.Context(), su.clusterId).
		Body(client.V1beta1ClusterServiceResetRootPasswordBody{RootPassword: su.password})
	_, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to reset root password: %v", err)
	}
	t.Logf("Root password reset successfully")
}

func (su *auditLogV2CustomerBucketTest) Test_04_EnableAuditLogToCustomer(t *testing.T) {
	var bucketURI string
	var updateReq client.ApiDatabaseAuditLogServiceUpdateAuditLogConfigRequest

	switch su.provider {
	case openapicommonv1beta1.Region_aws.String():
		bucketURI = fmt.Sprintf("s3://%s", su.awsS3BucketName)
		updateReq = su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
			AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
				Enabled:        *client.NewNullableBool(lo.ToPtr(true)),
				BucketManager:  lo.ToPtr(client.AUDITLOGCONFIGBUCKETMANAGER_CUSTOMER),
				BucketUri:      lo.ToPtr(bucketURI),
				AwsRoleArn:     lo.ToPtr(su.awsS3RoleArn),
				BucketRegionId: lo.ToPtr(fmt.Sprintf("%s-%s", su.provider, su.regionCode)),
			})
	case openapicommonv1beta1.Region_azure.String():
		bucketURI = fmt.Sprintf("https://%s.blob.core.windows.net/%s", su.azureBlobAccountName, su.azureBlobContainer)
		updateReq = su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
			AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
				Enabled:       *client.NewNullableBool(lo.ToPtr(true)),
				BucketManager: lo.ToPtr(client.AUDITLOGCONFIGBUCKETMANAGER_CUSTOMER),
				BucketUri:     lo.ToPtr(bucketURI),
				AzureSasToken: lo.ToPtr(su.azureBlobSASToken),
			})
	case openapicommonv1beta1.Region_gcp.String():
		bucketURI = fmt.Sprintf("gcs://%s", su.gcpGCSBucketName)
		updateReq = su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
			AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
				Enabled:        *client.NewNullableBool(lo.ToPtr(true)),
				BucketManager:  lo.ToPtr(client.AUDITLOGCONFIGBUCKETMANAGER_CUSTOMER),
				BucketUri:      lo.ToPtr(bucketURI),
				BucketRegionId: lo.ToPtr(fmt.Sprintf("%s-%s", su.provider, su.regionCode)),
			})
	default:
		t.Fatalf("unsupported provider: %s", su.provider)
	}

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update audit log config: %v", err)
	}
	if !resp.GetEnabled() {
		t.Fatal("expected audit log to be enabled")
	}
	if resp.GetBucketManager() != client.AUDITLOGCONFIGBUCKETMANAGER_CUSTOMER {
		t.Fatalf("expected bucket manager to be CUSTOMER, got %s", resp.GetBucketManager())
	}
	if resp.GetBucketUri() != bucketURI {
		t.Fatalf("expected bucket uri %s, got %s", bucketURI, resp.GetBucketUri())
	}

	// Wait for 3 minutes to make sure the audit log is enabled.
	enabled := false
	for range 18 {
		req := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
		resp, _, err := req.Execute()
		if err != nil {
			t.Fatalf("failed to get audit log config: %v", err)
		}
		if resp.GetEnabled() && resp.GetBucketManager() == client.AUDITLOGCONFIGBUCKETMANAGER_CUSTOMER {
			enabled = true
			break
		}
		time.Sleep(10 * time.Second)
	}
	if !enabled {
		t.Fatal("store audit log to customer failed")
	}

	// Sleep to ensure changes are applied
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_05_SetLogFormatToJSON(t *testing.T) {
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
		AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
			Format: lo.ToPtr(client.V1BETA1AUDITLOGCONFIGFORMAT_JSON),
		})

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update audit log format: %v", err)
	}
	if resp.GetFormat() != client.V1BETA1AUDITLOGCONFIGFORMAT_JSON {
		t.Fatalf("expected log format to be JSON, got %s", resp.GetFormat())
	}

	// Verify the format was updated
	req := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
	getResp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log config: %v", err)
	}
	if getResp.GetFormat() != client.V1BETA1AUDITLOGCONFIGFORMAT_JSON {
		t.Fatalf("expected log format to be JSON, got %s", getResp.GetFormat())
	}

	// Sleep to ensure changes are applied
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_06_ConfigureLogRotationBySize(t *testing.T) {
	// Set rotation size to 1MB
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
		AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
			RotationPolicy: &client.V1beta1AuditLogConfigRotationPolicy{
				MaxSizeMib: lo.ToPtr(int32(1)), // 1MB
			},
		})

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update log rotation size: %v", err)
	}
	if *resp.GetRotationPolicy().MaxSizeMib != 1 {
		t.Fatalf("expected rotation size to be 1MB, got %dMB", *resp.GetRotationPolicy().MaxSizeMib)
	}

	// Verify the rotation size was updated
	req := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
	getResp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log config: %v", err)
	}
	if *getResp.GetRotationPolicy().MaxSizeMib != 1 {
		t.Fatalf("expected rotation size to be 1MB, got %dMB", *getResp.GetRotationPolicy().MaxSizeMib)
	}

	// Sleep to ensure changes are applied
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_07_ConfigureLogRotationByTime(t *testing.T) {
	// Set rotation time to 1 minute
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
		AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
			RotationPolicy: &client.V1beta1AuditLogConfigRotationPolicy{
				MaxAgeSecond: *client.NewNullableInt32(lo.ToPtr(int32(60))), // 60 seconds
			},
		})

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update log rotation lifetime: %v", err)
	}
	if *resp.GetRotationPolicy().MaxAgeSecond.Get() != 60 {
		t.Fatalf(
			"expected rotation lifetime to be 60 seconds, got %d seconds",
			*resp.GetRotationPolicy().MaxAgeSecond.Get(),
		)
	}

	// Verify the rotation lifetime was updated
	req := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
	getResp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log config: %v", err)
	}
	if *getResp.GetRotationPolicy().MaxAgeSecond.Get() != 60 {
		t.Fatalf(
			"expected rotation lifetime to be 60 seconds, got %d seconds",
			*getResp.GetRotationPolicy().MaxAgeSecond.Get(),
		)
	}

	// Sleep to ensure changes are applied
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_08_DisableLogRedaction(t *testing.T) {
	// Disable log redaction
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
		AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
			Unredacted: *client.NewNullableBool(lo.ToPtr(true)),
		})

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update log redaction: %v", err)
	}
	if !resp.GetUnredacted() {
		t.Fatal("expected log redaction to be disabled")
	}

	// Verify the redaction was updated
	req := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
	getResp, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log config: %v", err)
	}
	if !getResp.GetUnredacted() {
		t.Fatal("expected log redaction to be disabled")
	}

	// Sleep to ensure changes are applied
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_09_CreateFilterRules(t *testing.T) {
	// Create first filter rule for all users with query events
	ruleReq1 := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceCreateAuditLogFilterRule(t.Context(), su.clusterId).
		AuditLogFilterRule(client.DatabaseAuditLogServiceCreateAuditLogFilterRuleRequest{
			UserExpr: lo.ToPtr("%@%"),
			EventTypes: []client.Dedicatedv1beta1AuditLogFilterRuleEventType{
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_QUERY,
			},
			DisplayName: lo.ToPtr("All-Users-Query-Rule"),
		})
	ruleResp1, _, err := ruleReq1.Execute()
	if err != nil {
		t.Fatalf("failed to create audit log filter rule: %v", err)
	}
	if ruleResp1.GetAuditLogFilterRuleId() == "" {
		t.Fatal("expected audit log filter rule id to be not empty")
	}
	su.filterRuleIds = append(su.filterRuleIds, ruleResp1.GetAuditLogFilterRuleId())

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)

	// Create second filter rule for root user with connection events
	ruleReq2 := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceCreateAuditLogFilterRule(t.Context(), su.clusterId).
		AuditLogFilterRule(client.DatabaseAuditLogServiceCreateAuditLogFilterRuleRequest{
			UserExpr: lo.ToPtr("root@%"),
			EventTypes: []client.Dedicatedv1beta1AuditLogFilterRuleEventType{
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_CONNECTION,
			},
			DisplayName: lo.ToPtr("Root-Connection-Rule"),
		})
	ruleResp2, _, err := ruleReq2.Execute()
	if err != nil {
		t.Fatalf("failed to create second audit log filter rule: %v", err)
	}
	if ruleResp2.GetAuditLogFilterRuleId() == "" {
		t.Fatal("expected second audit log filter rule id to be not empty")
	}
	su.filterRuleIds = append(su.filterRuleIds, ruleResp2.GetAuditLogFilterRuleId())

	// Verify rules were created
	listReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceListAuditLogFilterRules(t.Context(), su.clusterId).
		PageSize(pageSize)
	listResp, _, err := listReq.Execute()
	if err != nil {
		t.Fatalf("failed to list audit log filter rules: %v", err)
	}

	if len(listResp.GetAuditLogFilterRules()) < 2 {
		t.Fatalf("expected at least 2 filter rules, got %d", len(listResp.GetAuditLogFilterRules()))
	}

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_10_UpdateFilterRule(t *testing.T) {
	if len(su.filterRuleIds) == 0 {
		t.Fatal("no filter rule IDs available for update test")
	}

	ruleId := su.filterRuleIds[0]

	// Get the rule first
	getReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getResp, _, err := getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log filter rule: %v", err)
	}

	// Update the rule to include both QUERY and CONNECTION events
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceReplaceAuditLogFilterRule(t.Context(), su.clusterId, ruleId).
		AuditLogFilterRule(client.RequiredTheAuditLogFilterRuleToUpdate{
			UserExpr: getResp.UserExpr,
			EventTypes: []client.Dedicatedv1beta1AuditLogFilterRuleEventType{
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_QUERY,
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_CONNECTION,
			},
			DisplayName: lo.ToPtr(getResp.GetDisplayName() + "-Updated"),
		})

	_, _, err = updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to update audit log filter rule: %v", err)
	}

	// Verify the rule was updated
	getUpdatedReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getUpdatedResp, _, err := getUpdatedReq.Execute()
	if err != nil {
		t.Fatalf("failed to get updated audit log filter rule: %v", err)
	}

	if len(getUpdatedResp.GetEventTypes()) != 2 {
		t.Fatalf("expected 2 event types after update, got %d", len(getUpdatedResp.GetEventTypes()))
	}

	if !strings.HasSuffix(getUpdatedResp.GetDisplayName(), "-Updated") {
		t.Fatalf("expected rule name to end with '-Updated', got %s", getUpdatedResp.GetDisplayName())
	}

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_11_DisableFilterRule(t *testing.T) {
	if len(su.filterRuleIds) == 0 {
		t.Fatal("no filter rule IDs available for disable test")
	}

	ruleId := su.filterRuleIds[0]

	// Get the rule first
	getReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getResp, _, err := getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log filter rule: %v", err)
	}

	// Disable the rule
	disableReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceReplaceAuditLogFilterRule(t.Context(), su.clusterId, ruleId).
		AuditLogFilterRule(client.RequiredTheAuditLogFilterRuleToUpdate{
			Disabled: lo.ToPtr(true),
			UserExpr: getResp.UserExpr,
			EventTypes: []client.Dedicatedv1beta1AuditLogFilterRuleEventType{
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_QUERY,
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_CONNECTION,
			},
			DisplayName: lo.ToPtr(getResp.GetDisplayName() + "-Updated"),
		})

	disableResp, _, err := disableReq.Execute()
	if err != nil {
		t.Fatalf("failed to disable audit log filter rule: %v", err)
	}
	if !*disableResp.GetAuditLogFilterRule().Disabled {
		t.Fatal("expected rule to be disabled, but it's still enabled")
	}

	// Verify the rule was disabled
	getReq = su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getResp, _, err = getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log filter rule: %v", err)
	}

	if !getResp.GetDisabled() {
		t.Fatal("expected rule to be disabled, but it's still enabled")
	}

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_12_EnableFilterRule(t *testing.T) {
	if len(su.filterRuleIds) == 0 {
		t.Fatal("no filter rule IDs available for enable test")
	}

	ruleId := su.filterRuleIds[0]

	// Get the rule first
	getReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getResp, _, err := getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log filter rule: %v", err)
	}

	// Enable the rule
	enableReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceReplaceAuditLogFilterRule(t.Context(), su.clusterId, ruleId).
		AuditLogFilterRule(client.RequiredTheAuditLogFilterRuleToUpdate{
			Disabled: lo.ToPtr(false),
			UserExpr: getResp.UserExpr,
			EventTypes: []client.Dedicatedv1beta1AuditLogFilterRuleEventType{
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_QUERY,
				client.DEDICATEDV1BETA1AUDITLOGFILTERRULEEVENTTYPE_CONNECTION,
			},
			DisplayName: lo.ToPtr(getResp.GetDisplayName() + "-Updated"),
		})

	enableResp, _, err := enableReq.Execute()
	if err != nil {
		t.Fatalf("failed to enable audit log filter rule: %v", err)
	}
	if *enableResp.GetAuditLogFilterRule().Disabled {
		t.Fatal("expected rule to be enabled, but it's still disabled")
	}

	// Verify the rule was enabled
	getReq = su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	getResp, _, err = getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log filter rule: %v", err)
	}

	if getResp.GetDisabled() {
		t.Fatal("expected rule to be enabled, but it's still disabled")
	}

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_13_DeleteFilterRule(t *testing.T) {
	if len(su.filterRuleIds) < 2 {
		t.Skip("not enough filter rule IDs available for delete test")
	}

	// Delete the second rule
	ruleId := su.filterRuleIds[1]

	deleteReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceDeleteAuditLogFilterRule(
		t.Context(),
		su.clusterId,
		ruleId,
	)
	_, _, err := deleteReq.Execute()
	if err != nil {
		t.Fatalf("failed to delete audit log filter rule: %v", err)
	}

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)

	// Verify the rule was deleted
	listReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceListAuditLogFilterRules(t.Context(), su.clusterId).
		PageSize(pageSize)
	listResp, _, err := listReq.Execute()
	if err != nil {
		t.Fatalf("failed to list audit log filter rules: %v", err)
	}

	if lo.ContainsBy(listResp.GetAuditLogFilterRules(), func(rule client.Dedicatedv1beta1AuditLogFilterRule) bool {
		return rule.GetAuditLogFilterRuleId() == ruleId
	}) {
		t.Fatal("expected rule to be deleted, but it's still in the list")
	}

	// Remove the deleted rule ID from our tracking list
	su.filterRuleIds = lo.Filter(su.filterRuleIds, func(id string, _ int) bool {
		return id != ruleId
	})

	// Sleep to ensure the async task is finished
	time.Sleep(20 * time.Second)
}

func (su *auditLogV2CustomerBucketTest) Test_14_SetPublicNetworkAccess(t *testing.T) {
	// Get the cluster first.
	getClusterReq := su.cli.ClusterServiceAPI.ClusterServiceGetCluster(t.Context(), su.clusterId)
	cluster, _, err := getClusterReq.Execute()
	if err != nil {
		t.Fatalf("failed to get cluster: %v", err)
	}
	if len(cluster.TidbNodeSetting.GetTidbNodeGroups()) != 1 {
		t.Fatalf("expected 1 tidb node group, got %d", len(cluster.TidbNodeSetting.GetTidbNodeGroups()))
	}

	group := cluster.TidbNodeSetting.GetTidbNodeGroups()[0]

	nullableTrue := client.NullableBool{}
	nullableTrue.Set(lo.ToPtr(true))

	updatePublicAccessReq := su.cli.TidbNodeGroupServiceAPI.TidbNodeGroupServiceUpdatePublicEndpointSetting(t.Context(), su.clusterId, group.GetTidbNodeGroupId()).
		PublicEndpointSetting(client.TidbNodeGroupServiceUpdatePublicEndpointSettingRequest{
			Enabled: nullableTrue,
			IpAccessList: []client.V1beta1PublicEndpointSettingIpAccessList{
				{
					CidrNotation: "0.0.0.0/0",
					Description:  lo.ToPtr("allow access from anywhere."),
				},
			},
		})
	_, _, err = updatePublicAccessReq.Execute()
	if err != nil {
		t.Fatalf("failed to update public endpoint setting: %v", err)
	}
}

func (su *auditLogV2CustomerBucketTest) Test_15_ExecuteSql(t *testing.T) {
	for range 30 {
		time.Sleep(10 * time.Second)

		// List TiDB groups first.
		listGroupsReq := su.cli.TidbNodeGroupServiceAPI.TidbNodeGroupServiceListTidbNodeGroups(t.Context(), su.clusterId).
			PageSize(pageSize)
		listGroupsResp, _, err := listGroupsReq.Execute()
		if err != nil {
			t.Fatalf("failed to list tidb node groups: %v", err)
		}
		if len(listGroupsResp.GetTidbNodeGroups()) != 1 {
			t.Fatalf("expected 1 tidb node group, got %d", len(listGroupsResp.GetTidbNodeGroups()))
		}

		pubEndpoint, ok := lo.Find(
			listGroupsResp.GetTidbNodeGroups()[0].GetEndpoints(),
			func(item client.Dedicatedv1beta1TidbNodeGroupEndpoint) bool {
				return item.GetConnectionType() == client.V1BETA1TIDBNODEGROUPENDPOINTCONNECTIONTYPE_PUBLIC
			},
		)
		if !ok {
			t.Fatal("public endpoint not found")
		}

		// Skip if the public endpoint is not ready.
		if pubEndpoint.GetHost() == "" {
			continue
		}

		if err := sql.ExecuteSqlWithUrl(t.Context(), pubEndpoint.GetHost(), "", su.password, "skip-verify", su.port); err != nil {
			t.Logf("Failed to execute SQL: %v, retrying...", err)
			continue
		}
		return
	}

	t.Fatal("Failed to execute SQL after multiple attempts")
}

func (su *auditLogV2CustomerBucketTest) Test_16_PauseCluster(t *testing.T) {
	// Pause the cluster.
	req := su.cli.ClusterServiceAPI.ClusterServicePauseCluster(t.Context(), su.clusterId)
	_, _, err := req.Execute()
	if err != nil {
		t.Fatalf("failed to pause cluster: %v", err)
	}

	// Wait for a while to make sure the cluster is paused.
	paused := false
	for range 100 {
		req := su.cli.ClusterServiceAPI.ClusterServiceGetCluster(t.Context(), su.clusterId)
		resp, _, err := req.Execute()
		if err != nil {
			t.Fatalf("failed to get cluster: %v", err)
		}
		if resp.GetState() == client.COMMONV1BETA1CLUSTERSTATE_PAUSED {
			paused = true
			break
		}
		time.Sleep(10 * time.Second)
	}
	if !paused {
		t.Fatal("Pause cluster failed")
	}
}

func (su *auditLogV2CustomerBucketTest) Test_17_CheckAuditLogsInCustomerBucket(t *testing.T) {
	var existing bool
	var err error

	for range 30 {
		switch su.provider {
		case openapicommonv1beta1.Region_aws.String():
			existing, err = su.CheckAwsS3(t)
			if err != nil {
				t.Fatalf("failed to check aws s3: %v", err)
			}
		case openapicommonv1beta1.Region_azure.String():
			existing, err = su.CheckAzureBlob(t)
			if err != nil {
				t.Fatalf("failed to check azure blob: %v", err)
			}
		case openapicommonv1beta1.Region_gcp.String():
			existing, err = su.CheckGcpGcs(t)
			if err != nil {
				t.Fatalf("failed to check gcp gcs: %v", err)
			}
		default:
			t.Fatalf("unsupported provider: %s", su.provider)
		}

		if existing {
			break
		}
		time.Sleep(10 * time.Second)
	}
	if !existing {
		t.Fatal("No audit logs found in customer bucket")
	}
}

func (su *auditLogV2CustomerBucketTest) Test_18_DisableAuditLog(t *testing.T) {
	// Resume the cluster first since we can't modify a paused cluster
	resumeReq := su.cli.ClusterServiceAPI.ClusterServiceResumeCluster(t.Context(), su.clusterId)
	_, _, err := resumeReq.Execute()
	if err != nil {
		t.Fatalf("failed to resume cluster: %v", err)
	}

	// Wait for the cluster to be active
	active := false
	for range 100 {
		req := su.cli.ClusterServiceAPI.ClusterServiceGetCluster(t.Context(), su.clusterId)
		resp, _, err := req.Execute()
		if err != nil {
			t.Fatalf("failed to get cluster: %v", err)
		}
		if resp.GetState() == client.COMMONV1BETA1CLUSTERSTATE_ACTIVE {
			active = true
			break
		}
		time.Sleep(10 * time.Second)
	}
	if !active {
		t.Fatal("Resume cluster failed")
	}

	// Disable audit log
	updateReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceUpdateAuditLogConfig(t.Context(), su.clusterId).
		AuditLogConfig(client.DatabaseAuditLogServiceUpdateAuditLogConfigRequest{
			Enabled: *client.NewNullableBool(lo.ToPtr(false)),
		})

	resp, _, err := updateReq.Execute()
	if err != nil {
		t.Fatalf("failed to disable audit log: %v", err)
	}
	if resp.GetEnabled() {
		t.Fatal("expected audit log to be disabled")
	}

	// Verify audit log is disabled
	getReq := su.cli.DatabaseAuditLogServiceAPI.DatabaseAuditLogServiceGetAuditLogConfig(t.Context(), su.clusterId)
	getResp, _, err := getReq.Execute()
	if err != nil {
		t.Fatalf("failed to get audit log config: %v", err)
	}
	if getResp.GetEnabled() {
		t.Fatal("expected audit log to be disabled, but it's still enabled")
	}
}

func (su *auditLogV2CustomerBucketTest) CheckAwsS3(t *testing.T) (bool, error) {
	t.Helper()
	var opts []func(*config.LoadOptions) error
	if su.awsProfile != "" {
		opts = append(opts, config.WithSharedConfigProfile(su.awsProfile))
	}

	cfg, err := config.LoadDefaultConfig(t.Context(), opts...)
	if err != nil {
		return false, fmt.Errorf("unable to load SDK config, %w", err)
	}

	s3Client := s3.NewFromConfig(cfg)
	listObjectsInput := &s3.ListObjectsV2Input{
		Bucket: &su.awsS3BucketName,
		Prefix: &su.clusterId,
	}

	output, err := s3Client.ListObjectsV2(t.Context(), listObjectsInput)
	if err != nil {
		return false, err
	}

	return len(output.Contents) > 0, nil
}

func (su *auditLogV2CustomerBucketTest) CheckAzureBlob(t *testing.T) (bool, error) {
	t.Helper()
	containerURL := fmt.Sprintf(
		"https://%s.blob.core.windows.net/%s?%s",
		su.azureBlobAccountName,
		su.azureBlobContainer,
		su.azureBlobSASToken,
	)
	containerClient, err := container.NewClientWithNoCredential(containerURL, nil)
	if err != nil {
		return false, err
	}

	pager := containerClient.NewListBlobsFlatPager(&container.ListBlobsFlatOptions{
		Prefix: &su.clusterId,
	})

	if pager.More() {
		page, err := pager.NextPage(t.Context())
		if err != nil {
			if strings.Contains(err.Error(), "The specified container does not exist") {
				t.Logf("Azure container does not exist yet, retrying...")
				return false, nil
			}
			return false, err
		}
		if page.Segment != nil && len(page.Segment.BlobItems) > 0 {
			blobNames := lo.Map(page.Segment.BlobItems, func(item *container.BlobItem, _ int) string {
				return *item.Name
			})
			t.Logf("Found %d blobs in Azure container: %v", len(page.Segment.BlobItems), blobNames)
			return true, nil
		}
		t.Logf("No blobs found in Azure container for cluster %s, retrying...", su.clusterId)
	}
	return false, nil
}

func (su *auditLogV2CustomerBucketTest) CheckGcpGcs(t *testing.T) (bool, error) {
	t.Helper()

	var opts []option.ClientOption
	if su.gcpCredentialsFile != "" {
		opts = append(opts, option.WithCredentialsFile(su.gcpCredentialsFile))
	}

	gcsClient, err := storage.NewClient(t.Context(), opts...)
	if err != nil {
		return false, fmt.Errorf("storage.NewClient: %w", err)
	}
	defer gcsClient.Close()

	bucket := gcsClient.Bucket(su.gcpGCSBucketName)
	query := &storage.Query{Prefix: su.clusterId}
	it := bucket.Objects(t.Context(), query)

	_, err = it.Next()
	if errors.Is(err, iterator.Done) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("Bucket.Objects: %w", err)
	}

	return true, nil
}