package action

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"io"
	"testing"
	"time"

	"github.com/pingcap/kvproto/pkg/keyspacepb"
	meteringsdkcommon "github.com/pingcap/metering_sdk/common"
	meteringsdkconfig "github.com/pingcap/metering_sdk/config"
	meteringsdkstorage "github.com/pingcap/metering_sdk/storage"
	meteringsdkmetawriter "github.com/pingcap/metering_sdk/writer/meta"
	"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"
	"github.com/tidbcloud/cluster-service-ng/pkg/clustergateway"
	clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	mysqlschema "github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tikv/pd/server/apiv2/handlers"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type recordingObjectStorageProvider struct {
	uploads       []recordedUpload
	existingPaths map[string]bool
}

type recordedUpload struct {
	path string
	data []byte
}

var _ meteringsdkstorage.ObjectStorageProvider = (*recordingObjectStorageProvider)(nil)

func (p *recordingObjectStorageProvider) Upload(ctx context.Context, path string, data io.Reader) error {
	b, err := io.ReadAll(data)
	if err != nil {
		return err
	}
	p.uploads = append(p.uploads, recordedUpload{path: path, data: b})
	return nil
}

func (p *recordingObjectStorageProvider) Download(ctx context.Context, path string) (io.ReadCloser, error) {
	return nil, errors.New("not implemented")
}

func (p *recordingObjectStorageProvider) Delete(ctx context.Context, path string) error {
	return nil
}

func (p *recordingObjectStorageProvider) Exists(ctx context.Context, path string) (bool, error) {
	return p.existingPaths[path], nil
}

func (p *recordingObjectStorageProvider) List(ctx context.Context, prefix string) ([]string, error) {
	return nil, nil
}

func gunzipBytes(t *testing.T, b []byte) []byte {
	t.Helper()

	r, err := gzip.NewReader(bytes.NewReader(b))
	require.NoError(t, err)
	defer func() { _ = r.Close() }()

	out, err := io.ReadAll(r)
	require.NoError(t, err)
	return out
}

func TestParseKeyspaceEncryptionEnabled(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		input     string
		want      bool
		wantError bool
	}{
		{name: "empty", input: "", want: false},
		{name: "enabled true", input: `{"enabled":true,"cmek_id":"x"}`, want: true},
		{name: "enabled false", input: `{"enabled":false}`, want: false},
		{name: "missing enabled", input: `{"cmek_id":"x"}`, wantError: true},
		{name: "malformed", input: `{"enabled":`, wantError: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := parseKeyspaceEncryptionEnabled(tt.input)
			if tt.wantError {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tt.want, got)
		})
	}
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_IncludesStorageEncryptionEnabled(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
		CapacityMode: apiv1.CapacityMode_RCU_DRIVEN,
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{
			Config: map[string]string{"encryption": `{"enabled":true,"cmek_id":"x"}`},
		},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c1"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)
	unzipped := gunzipBytes(t, provider.uploads[0].data)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(unzipped, &meta))

	require.Equal(t, "premium", meta.Metadata["plan"])
	require.Equal(t, apiv1.CapacityMode_RCU_DRIVEN.String(), meta.Metadata["capacityMode"])
	v, ok := meta.Metadata["storageEncryptionEnabled"]
	require.True(t, ok)
	require.Equal(t, true, v)
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_UsesConfiguredPlan(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"byoc",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
		CapacityMode: apiv1.CapacityMode_RCU_DRIVEN,
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{
			Config: map[string]string{"encryption": `{"enabled":false}`},
		},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c1"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "byoc", meta.Metadata["plan"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_UsesEssentialPlanForSharedTiProxy(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}, DeployType: mysqlschema.SharedPoolDeployTypeSharedTiProxy}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{
			Config: map[string]string{"encryption": `{"enabled":false}`},
		},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c1"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "essential", meta.Metadata["plan"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_TopologyDrivenUsesPremiumReservedPlan(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"byoc",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c-topology",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
		CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN,
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}

	tidbClusterRepo.On("Load", mock.Anything, "c-topology", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c-topology").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{
			Config: map[string]string{"encryption": `{"enabled":false}`},
		},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c-topology"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
		utils.TraceIdKey:                      actions.ActionArgString("trace-topology-logic"),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "premium_reserved", meta.Metadata["plan"])
	require.Equal(t, apiv1.CapacityMode_TOPOLOGY_DRIVEN.String(), meta.Metadata["capacityMode"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_GetKeySpaceError(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(nil, errors.New("boom"))

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:  actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey: actions.ActionArgString("c1"),
	}

	err := action.Invoke(args)
	require.Error(t, err)
	require.Empty(t, provider.uploads)
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_DeleteIgnoresMissingKeyspace(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(nil, clustergateway.ErrKeyspaceNotFound)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c1"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(workflow.MeteringDeleteTiDBClusterValue),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, true, meta.Metadata["deleted"])
	require.Equal(t, false, meta.Metadata["storageEncryptionEnabled"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_TopologyDrivenWithActiveWorkerFetchesKeyspace(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c-topology-worker",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
		CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN,
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{
		Id:           "sp1",
		CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN,
		ComponentGroups: []*apiv1.Component{{
			ComponentType: apiv1.ComponentType_TIDB_WORKER,
			NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
		}},
	}}

	tidbClusterRepo.On("Load", mock.Anything, "c-topology-worker", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c-topology-worker").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{Config: map[string]string{"encryption": `{"enabled":true}`}},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c-topology-worker"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, true, meta.Metadata["storageEncryptionEnabled"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaByTiDBCluster_NilConfigReturnsFalse(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	tidbCluster := &apiv1.Cluster{
		Id:           "c1",
		TenantId:     "t1",
		SharedPoolId: "sp1",
		RequestUnit:  &apiv1.Cluster_RequestUnit{Min: 1, Max: 2},
	}
	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}

	tidbClusterRepo.On("Load", mock.Anything, "c1", true).Return(tidbCluster, nil)
	sharedPoolRepo.On("Load", mock.Anything, "sp1", true).Return(sharedPool, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{Config: nil},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:              actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster),
		workflow.TiDBClusterIdKey:             actions.ActionArgString("c1"),
		workflow.MeteringDeleteTiDBClusterKey: actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)

	unzipped := gunzipBytes(t, provider.uploads[0].data)
	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(unzipped, &meta))
	require.Equal(t, false, meta.Metadata["storageEncryptionEnabled"])
}

func TestSendMeteringMetadataAction_Invoke_LogicMetaBySharedPool_UsesClusterEncryptionFlag(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
	clusterOne := &apiv1.Cluster{Id: "c1", TenantId: "t1", SharedPoolId: "sp1", RequestUnit: &apiv1.Cluster_RequestUnit{Min: 1, Max: 2}}
	clusterTwo := &apiv1.Cluster{Id: "c2", TenantId: "t2", SharedPoolId: "sp1", RequestUnit: &apiv1.Cluster_RequestUnit{Min: 1, Max: 2}}

	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil)
	tidbClusterRepo.On("List", mock.Anything, &repo.TiDBClusterListFilter{SharedPoolIds: []string{"sp1"}}).Return([]*entity.TiDBCluster{clusterOne, clusterTwo}, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c1").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{Config: map[string]string{"encryption": `{"enabled":true}`}},
	}, nil)
	clusterGatewayCli.On("GetKeySpace", mock.Anything, sharedPool, "c2").Return(&handlers.KeyspaceMeta{
		KeyspaceMeta: &keyspacepb.KeyspaceMeta{Config: map[string]string{"encryption": ``}},
	}, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey: actions.ActionArgString(workflow.MeteringTypeValueLogicMetaBySharedPool),
		workflow.SharedPoolIdKey: actions.ActionArgString("sp1"),
		workflow.TenantIdKey:     actions.ActionArgString(""),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 2)

	firstMeta := meteringsdkcommon.MetaData{}
	secondMeta := meteringsdkcommon.MetaData{}
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &firstMeta))
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[1].data), &secondMeta))

	require.Equal(t, true, firstMeta.Metadata["storageEncryptionEnabled"])
	require.Equal(t, false, secondMeta.Metadata["storageEncryptionEnabled"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_DeleteTombstone(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN}}
	sharedPool.TenantId = "tenant-1"

	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:             actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey:             actions.ActionArgString("sp1"),
		workflow.TenantIdKey:                 actions.ActionArgString("tenant-1"),
		workflow.MeteringDeleteSharedPoolKey: actions.ActionArgString(workflow.MeteringDeleteSharedPoolValue),
		workflow.MeteringModifyTSKey:         actions.ActionArgString("1712793600"),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)
	unzipped := gunzipBytes(t, provider.uploads[0].data)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(unzipped, &meta))

	require.Equal(t, int64(1712793600), meta.ModifyTS)
	require.Equal(t, true, meta.Metadata["deleted"])
	require.Equal(t, "premium", meta.Metadata["plan"])
	require.Equal(t, apiv1.CapacityMode_RCU_DRIVEN.String(), meta.Metadata["capacityMode"])
	require.Equal(t, "sp1", meta.Metadata["physicalID"])
	require.Equal(t, "tenant-1", meta.Metadata["tenantId"])
}

func TestSendSharedPoolLifecycleMetadataAction(t *testing.T) {
	createdAt := time.Unix(1712793600, 0)
	tests := []struct {
		name         string
		actionName   string
		pauseState   string
		terminal     bool
		expectedTime int64
	}{
		{name: "pausing", actionName: SendPausingMeteringMetadataActionName, pauseState: SharedPoolMeteringStatePausing, expectedTime: createdAt.Unix()},
		{name: "paused", actionName: SendPausedMeteringMetadataActionName, pauseState: SharedPoolMeteringStatePaused, terminal: true},
		{name: "resuming", actionName: SendResumingMeteringMetadataActionName, pauseState: SharedPoolMeteringStateResuming, expectedTime: createdAt.Unix()},
		{name: "active", actionName: SendActiveMeteringMetadataActionName, pauseState: SharedPoolMeteringStateActive, terminal: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			operationRepo := repomnocks.NewPauseResumeOperationRepo(t)
			clusterGatewayCli := clustergatewaymocks.NewClient(t)
			provider := &recordingObjectStorageProvider{}
			writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
			metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)
			meteringAction := NewSendMeteringMetadataAction(
				tidbClusterRepo, sharedPoolRepo, clusterGatewayCli, metaWriter, "dev", "us-east-1", "", true,
			)
			lifecycleAction := NewSendSharedPoolLifecycleMetadataAction(
				tt.actionName, tt.pauseState, tt.terminal, operationRepo, meteringAction,
			)

			operation := &entity.PauseResumeOperation{
				SharedPoolId: "sp1",
				OperationId:  "op1",
				CreatedAt:    timestamppb.New(createdAt),
			}
			sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}
			sharedPool.TenantId = "tenant-1"
			operationRepo.On("LoadByID", mock.Anything, "sp1", "op1").Return(operation, nil).Once()
			sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil).Once()

			err := lifecycleAction.Invoke(actions.ActionArgs{
				workflow.SharedPoolIdKey:           actions.ActionArgString("sp1"),
				workflow.PauseResumeOperationIdKey: actions.ActionArgString("op1"),
			})
			require.NoError(t, err)
			require.Equal(t, tt.actionName, lifecycleAction.Name())
			require.Len(t, provider.uploads, 1)

			var meta meteringsdkcommon.MetaData
			require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
			if tt.terminal {
				require.Greater(t, meta.ModifyTS, createdAt.Unix())
				require.LessOrEqual(t, meta.ModifyTS, time.Now().Unix())
			} else {
				require.Equal(t, tt.expectedTime, meta.ModifyTS)
			}
			require.Equal(t, tt.pauseState, meta.Metadata["pauseState"])
			require.Equal(t, "op1", meta.Metadata["operationId"])
			require.Equal(t, false, meta.Metadata["deleted"])
		})
	}
}

func TestSendSharedPoolLifecycleMetadataAction_FencesCancelledOperation(t *testing.T) {
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	operationRepo := repomnocks.NewPauseResumeOperationRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	provider := &recordingObjectStorageProvider{}
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, nil)
	meteringAction := NewSendMeteringMetadataAction(
		repomnocks.NewTiDBClusterRepo(t), sharedPoolRepo, clustergatewaymocks.NewClient(t), metaWriter, "dev", "us-east-1", "", true,
	)
	lifecycleAction := NewSendSharedPoolLifecycleMetadataAction(
		SendPausingMeteringMetadataActionName, SharedPoolMeteringStatePausing, false, operationRepo, meteringAction,
	)
	guard := NewPauseResumeOwnerGuard(operationRepo, sharedPoolRepo, lockRepo)
	guardedAction := actions.GuardAction(lifecycleAction, guard.Run)

	lockRepo.EXPECT().Lock(mock.Anything, mock.Anything).Return(func() error { return nil }, nil).Once()
	operationRepo.EXPECT().LoadByID(mock.Anything, "sp1", "op1").Return(&entity.PauseResumeOperation{
		SharedPoolId: "sp1",
		OperationId:  "op1",
		Type:         apiv1.PauseResumeOperation_PAUSE,
		Status:       apiv1.PauseResumeOperation_CANCELLED,
	}, nil).Once()

	err := guardedAction.Invoke(actions.ActionArgs{
		workflow.SharedPoolIdKey:             actions.ActionArgString("sp1"),
		workflow.PauseResumeOperationIdKey:   actions.ActionArgString("op1"),
		workflow.PauseResumeOperationTypeKey: actions.ActionArgString(apiv1.PauseResumeOperation_PAUSE.String()),
	})
	require.Error(t, err)
	require.True(t, actions.IsTerminalError(err))
	require.Empty(t, provider.uploads)
}

func TestSendSharedPoolLifecycleMetadataAction_RetriesTimestampCollision(t *testing.T) {
	createdAt := time.Unix(1712793600, 0)
	collidingPath := "metering/meta/sharedpool/sp1/1712793600.json.gz"
	provider := &recordingObjectStorageProvider{existingPaths: map[string]bool{collidingPath: true}}
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, nil)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	operationRepo := repomnocks.NewPauseResumeOperationRepo(t)
	meteringAction := NewSendMeteringMetadataAction(
		repomnocks.NewTiDBClusterRepo(t), sharedPoolRepo, clustergatewaymocks.NewClient(t), metaWriter, "dev", "us-east-1", "", true,
	)
	lifecycleAction := NewSendSharedPoolLifecycleMetadataAction(
		SendResumingMeteringMetadataActionName, SharedPoolMeteringStateResuming, false, operationRepo, meteringAction,
	)
	operationRepo.On("LoadByID", mock.Anything, "sp1", "op1").Return(&entity.PauseResumeOperation{
		SharedPoolId: "sp1",
		OperationId:  "op1",
		CreatedAt:    timestamppb.New(createdAt),
	}, nil).Once()
	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{Id: "sp1"},
	}, nil).Once()

	err := lifecycleAction.Invoke(actions.ActionArgs{
		workflow.SharedPoolIdKey:           actions.ActionArgString("sp1"),
		workflow.PauseResumeOperationIdKey: actions.ActionArgString("op1"),
	})
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)
	require.NotEqual(t, collidingPath, provider.uploads[0].path)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Greater(t, meta.ModifyTS, createdAt.Unix())
	require.Equal(t, SharedPoolMeteringStateResuming, meta.Metadata["pauseState"])
	require.Equal(t, "op1", meta.Metadata["operationId"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_UsesConfiguredPlan(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"byoc",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}
	sharedPool.TenantId = "tenant-1"

	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey: actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey: actions.ActionArgString("sp1"),
		utils.TraceIdKey:         actions.ActionArgString("trace-topology-shared"),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "premium_reserved", meta.Metadata["plan"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_TopologyDrivenUsesPremiumReservedPlan(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"byoc",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN}}
	sharedPool.TenantId = "tenant-1"

	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey: actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey: actions.ActionArgString("sp1"),
	}

	err := action.Invoke(args)
	require.NoError(t, err)
	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "premium_reserved", meta.Metadata["plan"])
	require.Equal(t, apiv1.CapacityMode_TOPOLOGY_DRIVEN.String(), meta.Metadata["capacityMode"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_UsesEssentialPlanForSharedTiProxy(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{Id: "sp1"},
		DeployType: mysqlschema.SharedPoolDeployTypeSharedTiProxy,
	}
	sharedPool.TenantId = "tenant-1"

	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(sharedPool, nil)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey: actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey: actions.ActionArgString("sp1"),
	}

	err := action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(gunzipBytes(t, provider.uploads[0].data), &meta))
	require.Equal(t, "essential", meta.Metadata["plan"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_DeleteTombstone_UsesSharedPoolSnapshot(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}
	sharedPool.TenantId = "tenant-from-snapshot"
	sharedPoolJSON, err := protojson.Marshal(sharedPool)
	require.NoError(t, err)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:             actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey:             actions.ActionArgString("sp1"),
		workflow.SharedPoolKey:               actions.ActionArgString(string(sharedPoolJSON)),
		workflow.TenantIdKey:                 actions.ActionArgString("tenant-from-args"),
		workflow.MeteringDeleteSharedPoolKey: actions.ActionArgString(workflow.MeteringDeleteSharedPoolValue),
		workflow.MeteringModifyTSKey:         actions.ActionArgString("1712793660"),
	}

	err = action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)
	unzipped := gunzipBytes(t, provider.uploads[0].data)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(unzipped, &meta))

	require.Equal(t, int64(1712793660), meta.ModifyTS)
	require.Equal(t, true, meta.Metadata["deleted"])
	require.Equal(t, "sp1", meta.Metadata["physicalID"])
	require.Equal(t, "tenant-from-args", meta.Metadata["tenantId"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_NonDeleteStillUsesRepoLoad(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	loadedSharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}
	loadedSharedPool.TenantId = "tenant-from-repo"
	sharedPoolRepo.On("Load", mock.Anything, "sp1", false).Return(loadedSharedPool, nil).Once()

	snapshotPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}
	snapshotPool.TenantId = "tenant-from-snapshot"
	sharedPoolJSON, err := protojson.Marshal(snapshotPool)
	require.NoError(t, err)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey: actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey: actions.ActionArgString("sp1"),
		workflow.SharedPoolKey:   actions.ActionArgString(string(sharedPoolJSON)),
	}

	err = action.Invoke(args)
	require.NoError(t, err)

	require.Len(t, provider.uploads, 1)
	unzipped := gunzipBytes(t, provider.uploads[0].data)

	var meta meteringsdkcommon.MetaData
	require.NoError(t, json.Unmarshal(unzipped, &meta))
	require.Equal(t, false, meta.Metadata["deleted"])
	require.Equal(t, "tenant-from-repo", meta.Metadata["tenantId"])
}

func TestSendMeteringMetadataAction_Invoke_SharedPoolMeta_DeleteTombstone_InvalidModifyTS(t *testing.T) {
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	clusterGatewayCli := clustergatewaymocks.NewClient(t)

	provider := &recordingObjectStorageProvider{}
	writerCfg := meteringsdkconfig.DefaultConfig().WithOverwriteExisting(true)
	metaWriter := meteringsdkmetawriter.NewMetaWriter(provider, writerCfg)

	action := NewSendMeteringMetadataAction(
		tidbClusterRepo,
		sharedPoolRepo,
		clusterGatewayCli,
		metaWriter,
		"dev",
		"us-east-1",
		"",
		true,
	)

	sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "sp1"}}
	sharedPool.TenantId = "tenant-1"
	sharedPoolJSON, err := protojson.Marshal(sharedPool)
	require.NoError(t, err)

	args := actions.ActionArgs{
		workflow.MeteringTypeKey:             actions.ActionArgString(workflow.MeteringTypeValueSharedPoolMeta),
		workflow.SharedPoolIdKey:             actions.ActionArgString("sp1"),
		workflow.SharedPoolKey:               actions.ActionArgString(string(sharedPoolJSON)),
		workflow.MeteringDeleteSharedPoolKey: actions.ActionArgString(workflow.MeteringDeleteSharedPoolValue),
		workflow.MeteringModifyTSKey:         actions.ActionArgString("not-a-timestamp"),
	}

	err = action.Invoke(args)
	require.Error(t, err)
	require.Contains(t, err.Error(), "failed to parse metering modify ts")
	require.Empty(t, provider.uploads)
}

func TestIsDeleteSharedPoolMetering(t *testing.T) {
	t.Parallel()

	require.True(t, isDeleteSharedPoolMetering(workflow.MeteringDeleteSharedPoolValue))
	require.False(t, isDeleteSharedPoolMetering(""))
	require.False(t, isDeleteSharedPoolMetering("false"))
}