package br

import (
	"context"
	"fmt"
	"math"
	"time"

	v1alpha1 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/samber/lo"
	clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/cluster-service-ng/config"
	"github.com/tidbcloud/cluster-service-ng/pkg/clustergateway"
	"github.com/tidbcloud/cluster-service-ng/pkg/dataplane"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

/*
The challenges here are:
1. Multiple restores in one sharedpool, all of them may need to scale tikv, how to make the scale linearly without race condition?
2. One scalling tikv for one restore should not block others restore's scale tikv decision, or the restores can not be executed parellel.
Here is the algorithm to solve these problems:
For race condition:
1. Persist tikv scaling decision chain of a sharedpool in database.
2. Before making scale tikv decision, check if the chain is already in database. There are some challenges here:
    - The chain is not complete, some of the restores may already finished/canceled, so we need to build the chain based on the running restores.
    - We should find the running restores first, then find every node in the chain until we get the latest decision.
3. If Yes, build the chain, find last decision on the chain, and make decision based on the last decision.
4. If No, make decision based on the current tikv usage.
For parellel execution:
1. The scale tikv decision contains before scale spec of tikv and the target scale spec of tikv.
2. If current tikv ready nodes already exceed the target scale spec of tikv, we think the tikv scaling is already done for this restore even it may still scaling for others's restore.
For core data structure of the chain, you can refer to the ScaleTikvDecision field in the LightweightRestore schema.
*/

type ScaleManager interface {
	Scale(ctx context.Context, input *ScaleInput) error
	WaitAndSetScaleFinish(ctx context.Context, sharedPoolID string, scaleID uint) (done bool, err error)
	AnyFullfillingDecision(ctx context.Context, sharedPoolID string) (bool, error)
}

type ScaleInput struct {
	SharedPoolID     string
	BizType          string
	BizID            string
	FullfilDataBytes uint64
	PostScaleWriteDB func(ctx context.Context, scaleID uint) error
}

type scaleManagerImpl struct {
	sharedPoolRepo        repo.SharedPoolRepo
	scaleTikvDecisionRepo repo.ScaleTikvDecisionRepo
	clustergatewayCli     clustergateway.Client
	dataplaneManager      dataplane.Manager
	transactioner         db.Transactioner
	bizConfig             *config.BizConfig
}

func NewScaleManager(
	sharedPoolRepo repo.SharedPoolRepo,
	scaleTikvDecisionRepo repo.ScaleTikvDecisionRepo,
	clustergatewayCli clustergateway.Client,
	dataplaneManager dataplane.Manager,
	transactioner db.Transactioner,
	bizConfig *config.BizConfig,
) ScaleManager {
	return &scaleManagerImpl{
		sharedPoolRepo:        sharedPoolRepo,
		scaleTikvDecisionRepo: scaleTikvDecisionRepo,
		clustergatewayCli:     clustergatewayCli,
		dataplaneManager:      dataplaneManager,
		transactioner:         transactioner,
		bizConfig:             bizConfig,
	}
}

func (s *scaleManagerImpl) Scale(ctx context.Context, input *ScaleInput) error {
	logger := tilog.WithContextL(ctx).With(zap.String("sharedPoolID", input.SharedPoolID), zap.String("scaleBizType", input.BizType), zap.String("scaleBizID", input.BizID), zap.Uint64("scaleFullfilDataBytes", input.FullfilDataBytes))
	// make scale decision based on others decision first to increase success rate since collecting tikv usage during tikv scaling may fail
	basedDecision, err := s.getBasedDecision(ctx, input.SharedPoolID, time.Now())
	if err != nil {
		logger.Error("failed to get based scale decision", zap.Error(err))
		return fmt.Errorf("failed to get based scale decision: %w", err)
	}
	if basedDecision != nil {
		logger.Info("found based decision", zap.Any("basedDecision", basedDecision))
		err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
			return s.makeScaleDecision(ctx, input, &basedDecision.To, schema.BaseOnAnotherDecision, basedDecision.ID)
		})
		if err != nil {
			logger.Error("failed to make scale decision based on other decision", zap.Error(err), zap.Any("basedDecision", basedDecision))
			return fmt.Errorf("failed to make scale decision based on other decision: %w", err)
		}
		logger.Info("successfully made scale decision based on other decision", zap.Any("basedDecision", basedDecision))
		return nil
	}
	logger.Info("no based scale decision found at the first time, will try make scale decision based on tikv spec usage")
	// get tikv spec usage from dataplane
	// Notice: DO NOT move this logic to the middle of "based on other decision" and "based on tikv spec usage", since this step take time.
	//         keep the decision making and sent executed shortly, try best to avoid race condition. We don't want to implement a complicate
	//         distributed lock to orchestra two scaling decision in one sharedpool.
	tikvSpecUsage, err := s.getCurrentTiKVSpecUsage(ctx, input.SharedPoolID)
	if err != nil {
		logger.Error("failed to get tikv spec usage", zap.Error(err))
		return fmt.Errorf("failed to get tikv spec usage: %w", err)
	}
	logger.Info("successfully got tikv spec usage", zap.Any("tikvSpecUsage", tikvSpecUsage))

	logger.Info("try to make scale decision based on others decision again since collect tikv usage take times, there may be others decisions made during collecting tikv usage")
	basedDecision, err = s.getBasedDecision(ctx, input.SharedPoolID, time.Now())
	if err != nil {
		logger.Error("failed to get based scale decision at second time", zap.Error(err))
		return fmt.Errorf("failed to get based scale decision at second time: %w", err)
	}
	if basedDecision != nil {
		logger.Info("based scale decision found at second time", zap.Any("basedDecision", basedDecision))
		err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
			return s.makeScaleDecision(ctx, input, &basedDecision.To, schema.BaseOnAnotherDecision, basedDecision.ID)
		})
		if err != nil {
			logger.Error("failed to make scale decision based on other decision", zap.Error(err), zap.Any("basedDecision", basedDecision))
			return fmt.Errorf("failed to make scale decision based on other decision: %w", err)
		}
		logger.Info("successfully made scale decision based on other decision", zap.Any("basedDecision", basedDecision))
		return nil
	}

	logger.Info("no based scale decision found again, will make scale decision based on tikv spec usage")
	err = s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		return s.makeScaleDecision(ctx, input, tikvSpecUsage, schema.BaseOnTiKVUsage, 0)
	})
	if err != nil {
		logger.Error("failed to make scale decision based on tikv spec usage", zap.Error(err), zap.Any("tikvSpecUsage", tikvSpecUsage))
		return fmt.Errorf("failed to make scale decision based on tikv spec usage: %w", err)
	}
	logger.Info("successfully made scale decision based on tikv spec usage", zap.Any("tikvSpecUsage", tikvSpecUsage))
	return nil
}

func (s *scaleManagerImpl) WaitAndSetScaleFinish(ctx context.Context, sharedPoolID string, scaleID uint) (done bool, err error) {
	scaleDecision, err := s.scaleTikvDecisionRepo.Load(ctx, scaleID, true)
	if err != nil {
		return false, fmt.Errorf("failed to load scale decision: %w", err)
	}

	if scaleDecision.DeletedAt.Valid {
		tilog.Info(ctx, "scale decision is deleted, no need to wait for scale finish", zap.Uint("scaleID", scaleID))
		return true, nil
	}

	if scaleDecision.Status == schema.ScaleStatusReachExpected {
		tilog.Info(ctx, "scale decision is already reach expected", zap.Uint("scaleID", scaleID))
		return true, nil
	}

	tikvGroup, _, err := s.getTiKVGroupAndClient(ctx, sharedPoolID)
	if err != nil {
		return false, fmt.Errorf("failed to get tikv group and client: %w", err)
	}

	if uint32(tikvGroup.Status.ReadyReplicas) >= scaleDecision.To.NodeCount {
		tilog.Info(ctx, "tikv group is already exceeded or equal the target replicas, it's safe to feed data now",
			zap.Uint32("targetReplicas", scaleDecision.To.NodeCount),
			zap.Uint32("readyReplicas", uint32(tikvGroup.Status.ReadyReplicas)))

		scaleDecision.Status = schema.ScaleStatusReachExpected
		err = s.scaleTikvDecisionRepo.Save(ctx, scaleDecision)
		if err != nil {
			tilog.Error(ctx, "failed to set scale decision to expected", zap.Error(err), zap.Uint("scaleID", scaleID))
			return false, fmt.Errorf("failed to save scale decision: %w", err)
		}
		tilog.Info(ctx, "successfully set scale decision to expected", zap.Uint("scaleID", scaleID))
		return true, nil
	}
	tilog.Info(ctx, "tikv group is not exceeded the target replicas, wait for next check",
		zap.Uint32("targetReplicas", scaleDecision.To.NodeCount),
		zap.Uint32("readyReplicas", uint32(tikvGroup.Status.ReadyReplicas)))
	return false, nil
}

func (s *scaleManagerImpl) AnyFullfillingDecision(ctx context.Context, sharedPoolID string) (bool, error) {
	scaleDecisions, err := s.scaleTikvDecisionRepo.List(ctx, &repo.ScaleTikvDecisionFilter{
		SharedPoolID:   sharedPoolID,
		IncludeDeleted: false,
	})
	if err != nil {
		return false, fmt.Errorf("failed to list scale decisions: %w", err)
	}
	if len(scaleDecisions) > 0 {
		tilog.Info(ctx, "there are fullfilling scale decisions",
			zap.Any("sharedPoolID", sharedPoolID),
			zap.Any("scaleDecisions", scaleDecisions))
		return true, nil
	}
	tilog.Info(ctx, "there are no fullfilling scale decisions", zap.Any("sharedPoolID", sharedPoolID))
	return false, nil
}

// We need to build the decision chain, including early finished decisions(no feedding data any more), and make decision based on the last decision on the chain.
// It's a very hard algorithm since some of the decisions in chains may already done/deleted.
// Let's use a appoximate algorithm to solve this problem, it's not perfect, but it's good enough.
// 1. get two days scale decisions of the sharedpool, we assume a decision would not last more than two days.
// 2. build the decision chains of two days decisions, there may exist several decision chains since not all decisions fullfilling time are overlay.
// 3. find the chain with existing fullfilling decisions, and make decision based on the last decision on the chain
func (s *scaleManagerImpl) getBasedDecision(ctx context.Context, sharedPoolId string, now time.Time) (*schema.ScaleTikvDecision, error) {
	twoDaysAgo := now.Add(-2 * 24 * time.Hour)

	// the decision is executed, the biz is feeding data to tikv, not finished yet
	// if exists such decision, we can't make decision based on current tikv spec and usage
	executedFullfillingBizDecisions, err := s.scaleTikvDecisionRepo.List(ctx, &repo.ScaleTikvDecisionFilter{
		SharedPoolID:   sharedPoolId,
		IncludeDeleted: false,
		StartTime:      &twoDaysAgo,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to list other running scale decisions: %w", err)
	}

	if len(executedFullfillingBizDecisions) == 0 {
		return nil, nil //nolint:nilnil //returning nil,nil is intentional - no executed fulfilling decisions found
	}

	twoDaysScaleDecisions, err := s.scaleTikvDecisionRepo.List(ctx, &repo.ScaleTikvDecisionFilter{
		SharedPoolID:   sharedPoolId,
		IncludeDeleted: true,
		StartTime:      &twoDaysAgo,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to list two days scale decisions: %w", err)
	}

	lastestExecutedFullfillingBizDecision := getLatestDecision(twoDaysScaleDecisions)
	decisionChains := buildDecisionChains(twoDaysScaleDecisions)
	var theChain *scaleDecisionChain
	for _, decisionChain := range decisionChains {
		if decisionChain.ids[lastestExecutedFullfillingBizDecision.ID] {
			theChain = decisionChain
			break
		}
	}
	if theChain == nil {
		return nil, nil //nolint:nilnil //returning nil,nil is intentional - no decision chain found
	}
	return theChain.reversedOrder[0], nil
}

func getLatestDecision(decisions []*schema.ScaleTikvDecision) *schema.ScaleTikvDecision {
	latest := decisions[0]
	for _, decision := range decisions {
		if decision.CreatedAt.After(latest.CreatedAt) {
			latest = decision
		}
	}
	return latest
}

type scaleDecisionChain struct {
	// make it easy to find if a decision is in the chain
	ids map[uint]bool
	// ordered chains, the first element is last decision of the chain
	reversedOrder []*schema.ScaleTikvDecision
}

func buildDecisionChains(decisions []*schema.ScaleTikvDecision) []*scaleDecisionChain {
	// id to decision map
	decisionMap := make(map[uint]*schema.ScaleTikvDecision)
	for _, decision := range decisions {
		decisionMap[decision.ID] = decision
	}
	// build next decision map
	// if there is deicision chain: a->b->c, then nextDecision[a] = b, nextDecision[b] = c
	nextDecision := map[uint]uint{}
	for _, decision := range decisions {
		// the decision is not based on other decision, skip
		if decision.BaseOn != schema.BaseOnAnotherDecision {
			continue
		}
		// the based on decision not found in decisionMap, it pre-decision is made two days ago, skip
		if _, exist := decisionMap[decision.BaseID]; !exist {
			continue
		}
		nextDecision[decision.BaseID] = decision.ID
	}
	// build pre decision map
	// if there is deicision chain: a->b->c, then preDecision[b] = a, preDecision[c] = b
	preDecision := map[uint]uint{}
	for preDecisionId, nextDecisionId := range nextDecision {
		preDecision[nextDecisionId] = preDecisionId
	}

	// find the tail decisions of chains, which has no next decision.
	// for the chain a->b->c, the tail decision is c
	tailDecisions := make([]*schema.ScaleTikvDecision, 0)
	for _, decision := range decisions {
		if _, ok := nextDecision[decision.ID]; !ok {
			tailDecisions = append(tailDecisions, decision)
		}
	}

	decisionChains := make([]*scaleDecisionChain, 0)
	for _, decision := range tailDecisions {
		decisionChain := &scaleDecisionChain{
			ids:           map[uint]bool{decision.ID: true},
			reversedOrder: []*schema.ScaleTikvDecision{decision},
		}
		// build chain from tail to head
		curDecision := decision
		for {
			preDecisionId, exists := preDecision[curDecision.ID]
			if !exists {
				break
			}
			preDecision := decisionMap[preDecisionId]
			decisionChain.ids[preDecisionId] = true
			decisionChain.reversedOrder = append(decisionChain.reversedOrder, preDecision)
			curDecision = preDecision
		}
		decisionChains = append(decisionChains, decisionChain)
	}

	return decisionChains
}

func (s *scaleManagerImpl) getCurrentTiKVSpecUsage(ctx context.Context, sharedPoolId string) (*schema.TiKVSpecUsage, error) {
	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return nil, fmt.Errorf("failed to load target sharedPool: %w", err)
	}
	topology, err := s.clustergatewayCli.GetClusterTopology(ctx, sharedPool, &clusteragentv1.GetClusterTopologyReq{
		ClusterId:        sharedPoolId,
		IncludeTikvUsage: true,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to get cluster topology: %w", err)
	}
	if len(topology.TikvNodes) == 0 {
		return nil, fmt.Errorf("no tikv nodes found")
	}

	totalCapacityBytes := uint64(0)
	totalAvailableBytes := uint64(0)
	for _, node := range topology.TikvNodes {
		totalCapacityBytes += node.StorageCapacityBytes
		totalAvailableBytes += node.StorageAvailableBytes
	}
	if totalCapacityBytes == 0 {
		return nil, fmt.Errorf("total capacity bytes is 0")
	}
	usagePercentage := uint32(math.Ceil(100 - float64(totalAvailableBytes)*100/float64(totalCapacityBytes)))

	return &schema.TiKVSpecUsage{
		NodeCount:           uint32(len(topology.TikvNodes)),
		TotalCapacityBytes:  totalCapacityBytes,
		TotalAvailableBytes: totalAvailableBytes,
		UsagePercentage:     usagePercentage,
	}, nil
}

func (s *scaleManagerImpl) makeScaleDecision(ctx context.Context, scaleInput *ScaleInput, scaleFrom *schema.TiKVSpecUsage, baseOn string, baseId uint) error {
	scaleTo := calScaleDecisionTo(scaleInput.FullfilDataBytes, scaleFrom, s.bizConfig.DefaultTiBRConfig.TiKVUsagePercentageThreshold)
	tilog.Info(ctx, "successfully calculated scale decision", zap.Any("scaleFrom", scaleFrom), zap.Any("scaleTo", scaleTo))
	scaleDecision := assembleScaleDecision(scaleInput, scaleFrom, scaleTo, baseOn, baseId)
	tilog.Info(ctx, "successfully assembled scale decision", zap.Any("scaleDecision", scaleDecision))
	logger := tilog.WithContextL(ctx).With(zap.String("basedOn", baseOn), zap.Uint("basedId", baseId), zap.Any("scaleDecision", scaleDecision))

	err := s.scaleTikvDecisionRepo.Create(ctx, scaleDecision)
	if err != nil {
		tilog.Error(ctx, "failed to create scale decision", zap.Error(err), zap.Any("scaleDecision", scaleDecision))
		return fmt.Errorf("failed to create scale decision: %w", err)
	}
	tilog.Info(ctx, "successfully created scale decision", zap.Any("scaleDecision", scaleDecision))

	err = scaleInput.PostScaleWriteDB(ctx, scaleDecision.ID)
	if err != nil {
		tilog.Error(ctx, "failed to call PostScaleWriteDB", zap.Error(err), zap.Any("scaleDecision", scaleDecision))
		return fmt.Errorf("failed to call PostScaleWriteDB: %w", err)
	}

	err = s.sentScaleDecisionToOperator(ctx, scaleInput.SharedPoolID, scaleDecision)
	if err != nil {
		logger.Error("failed to sent scale decision to kernel", zap.Error(err))
		return fmt.Errorf("failed to sent scale decision to kernel: %w", err)
	}
	return nil
}

func assembleScaleDecision(scaleInput *ScaleInput, scaleFrom, scaleTo *schema.TiKVSpecUsage, baseOn string, baseId uint) *schema.ScaleTikvDecision {
	needScale := scaleTo.NodeCount > scaleFrom.NodeCount
	status := schema.ScaleStatusExecuting
	if !needScale {
		status = schema.ScaleStatusReachExpected
	}
	scaleDecision := &schema.ScaleTikvDecision{
		SharedPoolID: scaleInput.SharedPoolID,
		BizType:      scaleInput.BizType,
		BizID:        scaleInput.BizID,
		Status:       status,
		NeedScale:    needScale,
		BaseOn:       baseOn,
		BaseID:       baseId,
		From:         *scaleFrom,
		To:           *scaleTo,
	}
	return scaleDecision
}

const tikvNodeCountMultiple uint32 = 3

func calScaleDecisionTo(toFullfilDataBytes uint64, basedUsage *schema.TiKVSpecUsage, threshold uint32) *schema.TiKVSpecUsage {
	if threshold <= 0 || threshold >= 100 {
		threshold = DefaultTiKVUsagePercentageThreshold

	}
	totalUsageAfterFullfiled := basedUsage.TotalCapacityBytes - basedUsage.TotalAvailableBytes + toFullfilDataBytes
	totalUsagePercentageAfterFullfiled := uint32(math.Ceil(float64(totalUsageAfterFullfiled) * 100 / float64(basedUsage.TotalCapacityBytes)))
	// no need to scale
	if totalUsagePercentageAfterFullfiled <= threshold {
		return &schema.TiKVSpecUsage{
			NodeCount:           basedUsage.NodeCount,
			TotalCapacityBytes:  basedUsage.TotalCapacityBytes,
			TotalAvailableBytes: basedUsage.TotalAvailableBytes - toFullfilDataBytes,
			UsagePercentage:     totalUsagePercentageAfterFullfiled,
		}
	} else { // need to scale
		singleNodeCapacityBytes := basedUsage.TotalCapacityBytes / uint64(basedUsage.NodeCount)
		expectedNodeCount := uint32(math.Ceil(float64(totalUsageAfterFullfiled) / float64(singleNodeCapacityBytes) / (float64(threshold) / 100)))
		targetNodeCount := roundUpToMultiple(expectedNodeCount, tikvNodeCountMultiple)
		totalCapacityBytes := singleNodeCapacityBytes * uint64(targetNodeCount)
		return &schema.TiKVSpecUsage{
			NodeCount:           targetNodeCount,
			TotalCapacityBytes:  totalCapacityBytes,
			TotalAvailableBytes: totalCapacityBytes - totalUsageAfterFullfiled,
			UsagePercentage:     uint32(math.Ceil(float64(totalUsageAfterFullfiled) * 100 / float64(totalCapacityBytes))),
		}
	}
}

func roundUpToMultiple(value, multiple uint32) uint32 {
	if multiple == 0 {
		return value
	}
	remainder := value % multiple
	if remainder == 0 {
		return value
	}
	return value + multiple - remainder
}

func (s *scaleManagerImpl) sentScaleDecisionToOperator(ctx context.Context, sharedPoolId string, decision *schema.ScaleTikvDecision) error {
	tikvGroup, remoteCli, err := s.getTiKVGroupAndClient(ctx, sharedPoolId)
	if err != nil {
		return fmt.Errorf("failed to get tikv group and client: %w", err)
	}

	logger := tilog.WithContextL(ctx).With(
		zap.Any("sharedPool", sharedPoolId),
		zap.Uint32("targetReplicas", decision.To.NodeCount),
		zap.Uint32("currentReplicas", uint32(*tikvGroup.Spec.Replicas)),
		zap.Any("scaleDecision", decision))

	if uint32(*tikvGroup.Spec.Replicas) >= decision.To.NodeCount {
		logger.Info("replicas of tikv group is already exceeded the target replicas, no need to scale tikv")
		return nil
	}
	tikvGroup.Spec.Replicas = lo.ToPtr(int32(decision.To.NodeCount))

	err = remoteCli.Update(ctx, tikvGroup)
	if err != nil {
		logger.Error("failed to scale tikv group", zap.Error(err))
		return fmt.Errorf("failed to scale tikv group: %w", err)
	}

	logger.Info("successfully scaled tikv group")
	return nil
}

func (s *scaleManagerImpl) getTiKVGroupAndClient(ctx context.Context, sharedPoolId string) (*v1alpha1.TiKVGroup, client.Client, error) {
	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to load shared pool: %w", err)
	}
	remoteCli, err := s.dataplaneManager.GetClient(sharedPool.GetDataplaneName())
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get dataplane client: %w", err)
	}
	compId := sharedPool.GetDefaultComponenId()
	if compId == "" {
		return nil, nil, fmt.Errorf("default component id not found for shared pool %s", sharedPoolId)
	}
	tikvGroupName := fmt.Sprintf("db-%s", compId)
	var tikvGroup = &v1alpha1.TiKVGroup{}
	if err := remoteCli.Get(ctx, types.NamespacedName{Name: tikvGroupName, Namespace: sharedPool.GetDataPlaneNamespace()}, tikvGroup); err != nil {
		return nil, nil, fmt.Errorf("failed to get TiKVGroup %s in ns %s: %w", tikvGroupName, sharedPool.GetDataPlaneNamespace(), err)
	}
	return tikvGroup, remoteCli, nil
}