// Package worker add by blacknc <blacknc@163.com> at 2024/07/15 11:46:31
package worker

import (
	"context"
	"dataservices/commons/constants"
	"dataservices/commons/contextvar"
	"dataservices/commons/grpc"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"dataservices/ddb"
	"dataservices/internal/entity"
	"dataservices/internal/repository_ddb"
	centralSrv "dataservices/internal/service/central"
	service "dataservices/internal/service/common"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
	"github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	"github.com/tidbcloud/tiinf/errs"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
)

const (
	privateLinkFullModeInterval  = 30 * time.Minute
	privateLinkIncreModeInterval = 30 * time.Second
)

type privatelinkTaskWorker struct {
	ctx              context.Context
	fullModeLastTime time.Time
}

func NewPrivatelinkTaskWorker(ctx context.Context) Worker {
	var worker = privatelinkTaskWorker{
		ctx: ctx,
	}
	return &worker
}

func (worker *privatelinkTaskWorker) Name() string {
	return "privatelink_worker"
}

func (worker *privatelinkTaskWorker) RunningMsInterval() time.Duration {
	return privateLinkIncreModeInterval
}

func (worker *privatelinkTaskWorker) IsEnable() bool {
	return true
}

func (worker *privatelinkTaskWorker) Handle(tracer trace.Tracer) {
	ctx, span := tracer.Start(worker.ctx, "privatelinkworker")
	defer span.End()
	startTime := time.Now()
	logger.Info(ctx, "privatelink task worker beging...")
	defer func() {
		logger.Info(ctx, "privatelink task worker complete", zap.Duration("cost", time.Since(startTime)))
	}()

	var err error
	defer func() {
		ok := strconv.FormatBool(err == nil)
		metrics.WorkerTaskCounter.WithLabelValues("privatelink_controlloop", "", ok).Inc()
	}()

	orgMap, err := service.LoadDedicatedOrgWhitelist(ctx)
	if err != nil {
		logger.Error(ctx, "failed to load dedicated org whitelist", zap.Error(err))
		return
	}

	if time.Since(worker.fullModeLastTime) > privateLinkFullModeInterval {
		if err = worker.fullMode(ctx, orgMap); err != nil {
			logger.Error(ctx, "privatelink task worker failed with full mode", zap.Error(err))
		}
	} else {
		if err = worker.increMode(ctx, orgMap); err != nil {
			logger.Error(ctx, "privatelink task worker failed with incre mode", zap.Error(err))
		}
	}
}

func (worker *privatelinkTaskWorker) fullMode(ctx context.Context, orgMap map[uint64]struct{}) error {
	ctx = contextvar.WithLoggerField(ctx, zap.String("mode", "full mode"))
	startTime := time.Now()
	logger.Info(ctx, "privatelink task worker run with fullmode")
	service.SetDedicatedOrgWhitelist(orgMap)
	var orgIDs []uint64
	for orgID := range orgMap {
		orgIDs = append(orgIDs, orgID)
	}

	ddbCli := ddb.Client()

	// Traverse the org whitelist, load all relevant AWS dedicated clusters,
	// if the private link does not exist, create it.
	n := len(orgIDs)
	batchSize := 10
	allClusterMap := make(map[uint64]*core.Cluster, len(orgIDs))
	isDevTier := false
	for i := 0; i < n; i += batchSize {
		j := i + batchSize
		if j > n {
			j = n
		}
		logger.Info(ctx, "privatelink task worker list aws dedicated cluster by orgids", zap.Ints("range", []int{i, j}), zap.Uint64s("org_ids", orgIDs[i:j]))
		clusterMap, err := grpc.ListClusterByOrgs(ctx, orgIDs[i:j], &isDevTier, []string{constants.ClusterCloudProviderAWS})
		if err != nil {
			logger.Error(ctx, "failed to list clusters by orgids in privatelink task worker incre mode", zap.Uint64s("org_ids", orgIDs[i:j]), zap.Error(err))
			continue
		}

		clusterKeys := make([]repository_ddb.UserClusterKey, 0, len(clusterMap))
		for _, clusterInfo := range clusterMap {
			if !centralSrv.IsSupportShadowAccount(clusterInfo) {
				continue
			}
			allClusterMap[clusterInfo.Id] = clusterInfo
			clusterKeys = append(clusterKeys, repository_ddb.UserClusterKey{
				ProjectID: clusterInfo.ProjectId,
				ClusterID: clusterInfo.Id,
			})
		}
		if len(clusterKeys) == 0 {
			continue
		}
		userClusterList, err := repository_ddb.ListUserClusterByIDs(ctx, ddbCli, clusterKeys)
		if err != nil {
			logger.Error(ctx, "failed to list user cluster from ddb by cluster keys in private link task worker incre mode", zap.Error(err))
			continue
		}
		userClusterMap := utils.ConvertArrToMap(userClusterList, func(userCluster *entity.UserClusterInfo) uint64 { return userCluster.ClusterId })
		for clusterID := range clusterMap {
			clusterInfo := allClusterMap[clusterID]
			if clusterInfo == nil {
				continue
			}
			if userCluster, ok := userClusterMap[clusterID]; !ok || userCluster.EndpointId == "" || userCluster.PrivateDNS == "" {
				if !ok {
					logger.Warn(ctx, "user cluster not found in ddb", zap.Uint64("cluster_id", clusterID))
				}
				if err := service.PutPrivatelinkTask(ctx, clusterInfo.ProjectId, clusterID, clusterInfo.CloudProviderName, clusterInfo.RegionName, false, "worker_fullmode_whitelist"); err != nil {
					logger.Error(ctx, "failed to insert cluster privatelink task", zap.Uint64("cluster_id", clusterID), zap.Error(err))
				}
			}
		}
	}
	worker.fullModeLastTime = time.Now()
	logger.Info(ctx, "privatelink task worker check org dedicated cluster complete", zap.Duration("cost", time.Since(startTime)))

	// Load all dedicated clusters from DDB, and delete privatelink if the cluster status does not meet expectations
	var exclusiveStartKey map[string]types.AttributeValue
	var round, count int
	for {
		logger.Info(ctx, "privatelink task worker query dedicated cluster from ddb", zap.Int("round", round), zap.Int("count", count))
		lastEvaluatedKey, userClusterList, err := repository_ddb.QueryUserDedicatedClusterInfoWithPage(ctx, ddbCli, exclusiveStartKey, 20)
		if err != nil {
			logger.Error(ctx, "privatelink task worker query dedicated cluster from ddb failed", zap.Int("round", round), zap.Error(err))
			return err
		}
		for _, userCluster := range userClusterList {
			if _, ok := allClusterMap[userCluster.ClusterId]; ok {
				continue
			}
			clusterInfo, err := grpc.GetClusterWithoutPermCheck(ctx, userCluster.ClusterId)
			if err != nil && !errs.IsError(err, errs.ErrNotFound) {
				logger.Error(ctx, "failed to get cluster by id in private link task worker full mode", zap.Uint64("cluster_id", userCluster.ClusterId), zap.Error(err))
				continue
			}
			if err == nil {
				if _, ok := orgMap[clusterInfo.TenantId]; ok &&
					centralSrv.IsSupportShadowAccount(clusterInfo) ||
					userCluster.EndpointId == "" {
					continue
				}
			}
			isDeleteCluster := errs.IsError(err, errs.ErrNotFound)
			if err := service.PutPrivatelinkTask(ctx, userCluster.ProjectId, userCluster.ClusterId, userCluster.Cloud, userCluster.Region, isDeleteCluster, "worker_fullmode_ddb"); err != nil {
				logger.Error(ctx, "failed to insert cluster privatelink task", zap.Uint64("cluster_id", userCluster.ClusterId), zap.Error(err))
			}
		}
		if len(lastEvaluatedKey) == 0 {
			break
		}
		exclusiveStartKey = lastEvaluatedKey
		round++
		count += len(userClusterList)
	}
	return nil
}

func (worker *privatelinkTaskWorker) increMode(ctx context.Context, orgMap map[uint64]struct{}) error {
	ctx = contextvar.WithLoggerField(ctx, zap.String("mode", "incre mode"))
	logger.Info(ctx, "privatelink task worker run with incremode")
	orgIDs, err := service.GetDedicatedOrgWhitelistDiff(ctx, orgMap)
	if err != nil {
		return err
	}
	if len(orgIDs) == 0 {
		logger.Info(ctx, "privatelink task worker exit with org whitelist no changed")
		return nil
	}

	service.SetDedicatedOrgWhitelist(orgMap)

	n := len(orgIDs)
	batchSize := 10
	isDevTier := false
	for i := 0; i < n; i += batchSize {
		j := i + batchSize
		if j > n {
			j = n
		}
		logger.Info(ctx, "privatelink task worker list aws dedicated cluster by orgids", zap.Ints("range", []int{i, j}), zap.Uint64s("org_ids", orgIDs[i:j]))
		clusterMap, err := grpc.ListClusterByOrgs(ctx, orgIDs[i:j], &isDevTier, []string{constants.ClusterCloudProviderAWS})
		if err != nil {
			logger.Error(ctx, "failed to list clusters by orgids in privatelink task worker incre mode", zap.Uint64s("org_ids", orgIDs[i:j]), zap.Error(err))
			continue
		}

		for _, clusterInfo := range clusterMap {
			if err := service.PutPrivatelinkTask(ctx, clusterInfo.ProjectId, clusterInfo.Id, clusterInfo.CloudProviderName, clusterInfo.RegionName, false, "worker_incrmode"); err != nil {
				logger.Error(ctx, "failed to increase cluster privatelink task", zap.Uint64("org_id", clusterInfo.TenantId), zap.Uint64("cluster_id", clusterInfo.Id), zap.Error(err))
			}
		}
	}
	return nil
}