package metering

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/samber/lo"
	regionaldedicatedv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.uber.org/zap"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/regionserver"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	clusterrepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/cluster"
	billingrepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/metering"
	regionrepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/region"
	spurepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/pkg/config"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator"
)

const (
	measureIntervalCluster = 1 * time.Minute
)

var (
	// billedClusterStateSet define which cluster states should be charged.
	// clusters of PAUSING,PAUSED,RESUMING should be charged
	// clusters of CREATING, UNAVAILABLE, DELETED states should NOT be charged
	billedClusterStateSet = []globaldedicatedv1.Cluster_State{
		globaldedicatedv1.Cluster_AVAILABLE,
		globaldedicatedv1.Cluster_PAUSING,
		globaldedicatedv1.Cluster_PAUSED,
		globaldedicatedv1.Cluster_RESUMING,
		globaldedicatedv1.Cluster_SCALING,
		globaldedicatedv1.Cluster_MODIFYING,
	}

	// billedNodeStateSet define which node states should be charged.
	// nodes of CREATING state should NOT be charged
	billedNodeStateSet = map[regionaldedicatedv1.ComponentNodeInfo_Node_State]bool{
		regionaldedicatedv1.ComponentNodeInfo_Node_NORMAL:      true,
		regionaldedicatedv1.ComponentNodeInfo_Node_UNAVAILABLE: true,
		regionaldedicatedv1.ComponentNodeInfo_Node_LEAVING:     true,
		regionaldedicatedv1.ComponentNodeInfo_Node_TERMINATING: true,
	}

	tracer = otel.Tracer("metering")
)

type RecordMeteringSvc struct {
	clusterRepo             clusterrepo.Repo
	accountCli              account.Client
	clusterMeteringRecorder *clusterMeteringRecorder
	featGateCli             featuregate.Service
	cfg                     *config.Config
}

func NewRecordMeteringSvc(
	clusterRepo clusterrepo.Repo,
	accountCli account.Client,
	trans db.TransProvider,
	skuRepo spurepo.SkuRepo,
	clusterMeteringRepo billingrepo.ClusterMeteringRepo,
	billingReportRecordRepo billingrepo.ReportRecordRepo,
	regionRepo regionrepo.Repo,
	idGeneratorFacade idgenerator.IDGenerator,
	regionServerCli regionserver.Client,
	featGateCli featuregate.Service,
	cfg *config.Config,
) *RecordMeteringSvc {
	return &RecordMeteringSvc{
		clusterRepo: clusterRepo,
		accountCli:  accountCli,
		clusterMeteringRecorder: &clusterMeteringRecorder{
			trans:                   trans,
			clusterRepo:             clusterRepo,
			skuRepo:                 skuRepo,
			clusterMeteringRepo:     clusterMeteringRepo,
			billingReportRecordRepo: billingReportRecordRepo,
			regionRepo:              regionRepo,
			idGeneratorFacade:       idGeneratorFacade,
			regionServerCli:         regionServerCli,
			featGateCli:             featGateCli,
		},
		featGateCli: featGateCli,
		cfg:         cfg,
	}
}

// RunClusterMetering is the entrypoint to measure cluster metering for all
func (svc *RecordMeteringSvc) RunClusterMetering(ctx context.Context, now time.Time) (err bizerr.BizError) {
	logger := tilog.WithContextS(ctx).With(zap.Time("now", now))
	ctx, span := tracer.Start(ctx, "RecordMeteringSvc.RunClusterMetering")
	defer span.End()
	_startAt := time.Now()
	defer func() {
		execDuration := time.Since(_startAt)
		success := "true"
		if err != nil {
			success = "false"
		}

		labels := map[string]string{
			MetricsLblTenantId:    "0",
			MetricsLblProjectId:   "0",
			MetricsLblClusterId:   "0",
			metricsLabelMeterType: metricsMeterTypeCluster,
			metricsLabelSuccess:   success,
		}
		MetricsMeterCount.With(labels).Inc()
		MetricsMeterDuration.With(labels).Observe(execDuration.Seconds())
	}()

	clusterMeteringFeatureGate, ok := svc.cfg.ServerConfig.FeatureGates[featuregate.ClusterMetering]
	if ok {
		logger.Info("cluster_metering feature gate config",
			zap.Any("DisabledTenants", clusterMeteringFeatureGate.DisabledTenants),
			zap.Any("DisabledProjects", clusterMeteringFeatureGate.DisabledProjects),
		)
	}

	clusters, err := svc.clusterRepo.Query(ctx, &clusterrepo.Filter{
		States: billedClusterStateSet,
	})
	if err != nil {
		return err.Chain("query clusters")
	}
	logger.Infof("there are %d clusters to meter", len(clusters))
	span.AddEvent("QueryClusters")

	projectIDSet := make(map[uint64]bool, len(clusters))
	projectIDs := make([]uint64, 0, len(clusters))
	for _, cluster := range clusters {
		_, ok := projectIDSet[cluster.ProjectId]
		if !ok {
			projectIDSet[cluster.ProjectId] = true
			projectIDs = append(projectIDs, cluster.ProjectId)
		}
	}
	id2Tenant, id2Project, err := svc.getProjectsAndTenants(ctx, projectIDs)
	if err != nil {
		return err.Chain("parse tenants and projects")
	}

	// We should do measure for projects by projectIDSet rather than using id2Project, because the id2Project may miss some projectID
	// if no project is associated with it.
	logger.Infof("there are %d projects to meter", len(projectIDSet))
	for projectID := range projectIDSet {
		project, ok := id2Project[projectID]
		if !ok {
			logger.Errorf("can't find project by id, projectId=%d", projectID)
			continue
		}
		tenant, ok := id2Tenant[project.TenantID]
		if !ok {
			logger.Errorf("can't find tenant by id, tenantId=%d, projectId=%d", project.TenantID, projectID)
			continue
		}

		if !svc.featGateCli.CheckFeatureEnabled(
			featuregate.ClusterMetering,
			project.TenantID,
			projectID,
			globalcommonv1.CloudProvider_NAME_UNSPECIFIED.String(),
			globalcommonv1.Region_NAME_UNSPECIFIED.String(),
		) {
			logger.Info("cluster_metering feature gate is disabled, skip", zap.Uint64("projectID", projectID))
			continue
		}

		if err = svc.clusterMeteringRecorder.Record(ctx, now, tenant, project); err != nil {
			logger.Warn(
				"failed to record cluster metering for project, skip",
				zap.Error(err),
				zap.Uint64("projectID", projectID),
			)
		}
	}
	return nil
}

func (svc *RecordMeteringSvc) getProjectsAndTenants(ctx context.Context, projectIDs []uint64) (
	map[uint64]*account.TenantInfo, map[uint64]*account.ProjectInfo, bizerr.BizError,
) {
	ctx, span := tracer.Start(ctx, "RecordMeteringSvc.getProjectsAndTenants")
	defer span.End()

	if len(projectIDs) == 0 {
		return nil, nil, nil
	}

	projectFilter := account.ProjectFilter{
		ProjectIDs: projectIDs,
	}
	id2Project, err := svc.accountCli.QueryProject(ctx, projectFilter)
	if err != nil {
		return nil, nil, err
	}

	tenantIDs := make([]uint64, 0, len(id2Project))
	for _, project := range id2Project {
		tenantIDs = append(tenantIDs, project.TenantID)
	}
	tenantFilter := account.TenantFilter{
		TenantIDs: tenantIDs,
	}
	id2Tenant, err := svc.accountCli.QueryTenant(ctx, tenantFilter)
	if err != nil {
		return nil, nil, err
	}

	return id2Tenant, id2Project, nil
}

type clusterMeteringRecorder struct {
	trans                   db.TransProvider
	clusterRepo             clusterrepo.Repo
	skuRepo                 spurepo.SkuRepo
	clusterMeteringRepo     billingrepo.ClusterMeteringRepo
	billingReportRecordRepo billingrepo.ReportRecordRepo
	regionRepo              regionrepo.Repo
	idGeneratorFacade       idgenerator.IDGenerator
	regionServerCli         regionserver.Client
	featGateCli             featuregate.Service
}

func (svc *clusterMeteringRecorder) Record(
	ctx context.Context,
	now time.Time,
	tenant *account.TenantInfo,
	project *account.ProjectInfo,
) (err bizerr.BizError) {
	logger := tilog.WithContextS(ctx).With(zap.Uint64("projectID", project.ID), zap.Uint64("tenantID", tenant.ID))
	ctx, span := tracer.Start(ctx, "ClusterMetering")
	defer span.End()
	logger.Infof("try cluster metering for project (ignore azure clusters)")
	_startAt := time.Now()
	defer func() {
		_tenantID := strconv.FormatUint(tenant.ID, 10)
		_projectID := strconv.FormatUint(project.ID, 10)
		execDuration := time.Since(_startAt)
		success := "true"
		if err != nil {
			success = "false"
		}

		labels := map[string]string{
			MetricsLblTenantId:    _tenantID,
			MetricsLblProjectId:   _projectID,
			MetricsLblClusterId:   "0",
			metricsLabelMeterType: metricsMeterTypeCluster,
			metricsLabelSuccess:   success,
		}
		MetricsMeterCount.With(labels).Inc()
		MetricsMeterDuration.With(labels).Observe(execDuration.Seconds())
	}()

	clusters, err := svc.clusterRepo.QueryAndLoadComponents(ctx, &clusterrepo.Filter{
		OrganizationIds: []uint64{tenant.ID},
		ProjectIds:      []uint64{project.ID},
		States:          billedClusterStateSet,
		CloudProviderNames: []globalcommonv1.CloudProvider_Name{
			globalcommonv1.CloudProvider_AWS,
			globalcommonv1.CloudProvider_GCP,
		},
	})
	if err != nil {
		return err.Chain("query clusters")
	}
	logger.Infof("len(clusters)=%d", len(clusters))
	span.AddEvent("QueryClusters (ignore azure clusters)")

	// adjust time, floor to the nearest minute. all cost data are recorded in minute.
	endTime := now.Truncate(measureIntervalCluster)
	for _, cluster := range clusters {
		logger = logger.With(zap.Uint64("clusterID", cluster.ClusterId))

		if err := svc.trans.Transaction(ctx, func(ctx context.Context) error {
			return svc.recordOneCluster(ctx, endTime, tenant, project, cluster)
		}); err != nil {
			logger.Warnf("failed to build cluster metering and billing report record: %s", err)
			continue
		}
	}
	return nil
}

func (svc *clusterMeteringRecorder) recordOneCluster(
	ctx context.Context,
	now time.Time,
	tenant *account.TenantInfo,
	project *account.ProjectInfo,
	cluster *entity.Cluster,
) (err bizerr.BizError) {
	logger := tilog.WithContextS(ctx).With(zap.Uint64("clusterID", cluster.ClusterId))
	ctx, span := tracer.Start(ctx, "ClusterMetering.recordOneCluster")
	defer span.End()
	span.SetAttributes(attribute.String("cluster_id", fmt.Sprint(cluster.ClusterId)))
	_startAt := time.Now()
	defer func() {
		_tenantID := strconv.FormatUint(tenant.ID, 10)
		_projectID := strconv.FormatUint(project.ID, 10)
		_clusterID := strconv.FormatUint(cluster.ClusterId, 10)
		execDuration := time.Since(_startAt)
		success := "true"
		if err != nil {
			success = "false"
		}

		labels := map[string]string{
			MetricsLblTenantId:    _tenantID,
			MetricsLblProjectId:   _projectID,
			MetricsLblClusterId:   _clusterID,
			metricsLabelMeterType: metricsMeterTypeCluster,
			metricsLabelSuccess:   success,
		}
		MetricsMeterCount.With(labels).Inc()
		MetricsMeterDuration.With(labels).Observe(execDuration.Seconds())
	}()

	newMetering, err := svc.buildClusterMetering(ctx, now, cluster)
	if err != nil {
		if err.RawErrCode() == bizerr.ErrCodeOptimisticLockConflict {
			logger.Infof("optimistic lock conflict, skip. err: %s", err)
			return nil
		}
		return err.Chain("build cluster metering")
	}
	if !(newMetering.EndTs).After(newMetering.StartTS) {
		logger.Infof(
			"newMetering.EndTs=%s is NOT after newMetering.StartTS=%s, skip",
			newMetering.EndTs,
			newMetering.StartTS,
		)
		return bizerr.New(bizerr.ErrCodeOptimisticLockConflict, "wait next time")
	}
	span.AddEvent("build cluster metering")

	clusterRecord, err := svc.buildReportRecord(ctx, tenant, project, cluster, newMetering)
	if err != nil {
		return err.Chain("build cluster report records")
	}
	span.AddEvent("build cluster report record")

	if err := svc.clusterMeteringRepo.Save(ctx, newMetering); err != nil {
		return err.Chain("save cluster metering")
	}
	clusterRecord.ClusterMeteringID = newMetering.ID
	if err := svc.billingReportRecordRepo.Save(ctx, clusterRecord); err != nil {
		return err.Chain("save cluster report record")
	}
	span.AddEvent("save to db")
	return nil
}

func (svc *clusterMeteringRecorder) buildClusterMetering(
	ctx context.Context,
	now time.Time,
	cluster *entity.Cluster,
) (*entity.ClusterMetering, bizerr.BizError) {
	latestMetering, err := svc.clusterMeteringRepo.LoadAndLockLatestOne(ctx, cluster.ClusterId)
	if err != nil && err.RawErrCode() != bizerr.ErrCodeRecordNotFound {
		return nil, err.Chain("load and lock cluster metering, clusterID=%d", cluster.ClusterId)
	}
	if latestMetering == nil {
		// generate fake one
		latestMetering = &entity.ClusterMetering{
			ID:        0,
			TenantID:  cluster.OrganizationId,
			ProjectID: cluster.ProjectId,
			ClusterID: cluster.ClusterId,
			StartTS:   time.Unix(0, 0),
			EndTs:     now.Add(-measureIntervalCluster),
			Latest:    true,
		}
	}

	newMetering := &entity.ClusterMetering{
		ID:        0,
		TenantID:  cluster.OrganizationId,
		ProjectID: cluster.ProjectId,
		ClusterID: cluster.ClusterId,
		StartTS:   latestMetering.EndTs,
		EndTs:     now,
		Latest:    true,
	}

	return newMetering, nil
}

func (svc *clusterMeteringRecorder) buildReportRecord(
	ctx context.Context,
	tenant *account.TenantInfo,
	project *account.ProjectInfo,
	cluster *entity.Cluster,
	metering *entity.ClusterMetering,
) (*entity.BillingReportRecord, bizerr.BizError) {
	logger := tilog.WithContextS(ctx).With(zap.Uint64("clusterId", cluster.ClusterId))
	duration := metering.EndTs.Sub(metering.StartTS)
	record := &entity.BillingReportRecord{
		Data:              []*entity.BillingMeteringReport{},
		Status:            entity.BillingRecordStatusUnsynced,
		ClusterMeteringID: metering.ID,
		TenantID:          tenant.ID,
	}
	// don't charge if cluster is creating.
	if cluster.State == globaldedicatedv1.Cluster_CREATING {
		return record, nil
	}

	// merge all tidb and tiproxy components into one.
	var (
		tidbCount, tiproxyCount         uint32
		tidbComponent, tiproxyComponent *globaldedicatedv1.Cluster_ClusterComponent
		newComponents                   []*globaldedicatedv1.Cluster_ClusterComponent
	)
	for _, each := range cluster.ClusterComponents {
		if _, ok := entity.ComponentsHaveExtraGroup[each.GetComponentType()]; !ok {
			newComponents = append(newComponents, each)
			continue
		}
		if each.GetComponentType() == globaldedicatedv1.ComponentType_TIDB {
			tidbCount += each.GetCount()
			tidbComponent = each
		}
		if each.GetComponentType() == globaldedicatedv1.ComponentType_TIPROXY {
			tiproxyCount += each.GetCount()
			tiproxyComponent = each
		}
	}
	if tidbComponent != nil {
		tidbComponent.Count = lo.EmptyableToPtr(tidbCount)
		newComponents = append(newComponents, tidbComponent)
	}
	if tiproxyComponent != nil {
		tiproxyComponent.Count = lo.EmptyableToPtr(tiproxyCount)
		newComponents = append(newComponents, tiproxyComponent)
	}

	// get componentNodeInfos since we need to check component node state if cluster is scaling
	var (
		componentNodeInfos []*regionaldedicatedv1.ComponentNodeInfo
		err                bizerr.BizError
	)
	if cluster.State == globaldedicatedv1.Cluster_SCALING {
		componentNodeInfos, err = svc.getComponentNodeInfos(ctx, cluster)
		if err != nil {
			return nil, err
		}
	}

	for _, each := range newComponents {
		billedNodeQuantity := svc.getBilledNodeQuantity(ctx, cluster, componentNodeInfos, each)
		logger.Infof(
			"nodeQuantity, componentType=%s, clusterStatus=%s, in spec: %d, billed: %d",
			each.GetComponentType(),
			cluster.State.String(),
			each.GetCount(),
			billedNodeQuantity,
		)

		dataStorage := each.GetStorages()[globaldedicatedv1.Sku_DATA.String()]
		labels := buildBillingMeasureLabel(
			tenant,
			project,
			cluster.GetClusterId(),
			each.GetInstanceCpuArch(),
			each.GetCpuNum(),
			each.GetMemorySizeGi(),
			dataStorage.GetIops(),
			dataStorage.GetThroughputMibps(),
			each.GetComponentType(),
		)

		usages := buildBillingReportForComponent(
			svc.idGeneratorFacade,
			metering.TenantID,
			cluster.RegionId.String(),
			billedNodeQuantity,
			cluster.State,
			duration,
			each,
			labels,
			false,
			each.GetComponentType(),
			cluster.GetCloudProviderId(),
		)
		for _, usage := range usages {
			record.Data = append(record.Data, &entity.BillingMeteringReport{
				CostID:        usage.CostID,
				SKU:           usage.Sku,
				CostTimeStamp: metering.StartTS.Unix(),
				Cost:          usage.Cost,
				Labels:        labels,
				Memo:          fmt.Sprintf("componentType:%s,componentId:%d", each.GetComponentType(), each.GetId()),
			})
		}
	}

	return record, nil
}

func (svc *clusterMeteringRecorder) getBilledNodeQuantity(
	ctx context.Context,
	cluster *entity.Cluster,
	componentNodeInfos []*regionaldedicatedv1.ComponentNodeInfo,
	component *globaldedicatedv1.Cluster_ClusterComponent,
) int {
	logger := tilog.WithContextS(ctx).With(zap.Uint64("clusterId", cluster.ClusterId))
	logger.Info(
		"getBilledNodeQuantity",
		zap.String("componentType", component.GetComponentType().String()),
		zap.String("clusterState", cluster.State.String()),
	)

	// NOTE: only scaling cluster need to get billed node quantity from GetComponentNodeInfo
	if cluster.State != globaldedicatedv1.Cluster_SCALING {
		return int(component.GetCount())
	}

	if componentNodeInfos == nil {
		return 0
	}

	count := 0
	for _, componentNodeInfo := range componentNodeInfos {
		if componentTypeRegionalToGlobal(componentNodeInfo.ComponentType) != component.GetComponentType() {
			continue
		}
		for _, node := range componentNodeInfo.Nodes {
			if _, ok := billedNodeStateSet[node.State]; ok {
				count++
			}
		}
	}

	return count
}

func componentTypeRegionalToGlobal(componentType regionaldedicatedv1.ComponentType) globaldedicatedv1.ComponentType {
	switch componentType {
	case regionaldedicatedv1.ComponentType_TIKV:
		return globaldedicatedv1.ComponentType_TIKV
	case regionaldedicatedv1.ComponentType_TIDB:
		return globaldedicatedv1.ComponentType_TIDB
	case regionaldedicatedv1.ComponentType_TIFLASH:
		return globaldedicatedv1.ComponentType_TIFLASH
	case regionaldedicatedv1.ComponentType_TIFLASH_LEARNER:
		return globaldedicatedv1.ComponentType_TIFLASH_LEARNER
	case regionaldedicatedv1.ComponentType_PD:
		return globaldedicatedv1.ComponentType_PD
	case regionaldedicatedv1.ComponentType_TIPROXY:
		return globaldedicatedv1.ComponentType_TIPROXY
	default:
		return globaldedicatedv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED
	}
}

func (svc *clusterMeteringRecorder) getComponentNodeInfos(
	ctx context.Context,
	cluster *entity.Cluster,
) ([]*regionaldedicatedv1.ComponentNodeInfo, bizerr.BizError) {
	region, err := svc.regionRepo.SelectRegionById(ctx, cluster.RegionId, cluster.OrganizationId)
	if err != nil {
		return nil, err
	}

	getComponentNodeInfoRequest := &regionaldedicatedv1.GetComponentNodeInfosRequest{
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
	}
	var getComponentNodeInfoResponse *regionaldedicatedv1.GetComponentNodeInfosResponse
	getComponentNodeInfoResponse, err = svc.regionServerCli.GetComponentNodeInfo(
		ctx,
		region,
		getComponentNodeInfoRequest,
	)
	if err != nil {
		return nil, err
	}
	return getComponentNodeInfoResponse.GetComponentNodeInfos(), nil
}