package server

import (
	"context"
	"fmt"
	"strings"

	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/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

func (s *serverImpl) UpdateTiDBCluster(ctx context.Context, tidbCluster *entity.TiDBCluster, force bool) (*entity.TiDBCluster, error) {
	if tidbCluster == nil || tidbCluster.Id == "" {
		return nil, status.Errorf(codes.InvalidArgument, "tidbCluster or tidbCluster id is required")
	}

	if tidbCluster.GetRequestUnit() == nil {
		return nil, status.Errorf(codes.InvalidArgument, "tidbCluster request unit is required")
	}

	var dagExecId string
	var dagSubmitted bool
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		// load current tidb cluster
		currTiDBCluster, err := s.tidbClusterRepo.Load(ctx, tidbCluster.Id, false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load tidbCluster: %s", err)
			}
			// not found, check if it's in creating
			return status.Errorf(codes.NotFound, "tidb cluster %s not found", tidbCluster.Id)
		}

		if currTiDBCluster.GetPhase() == apiv1.Cluster_UNAVAILABLE {
			return status.Errorf(codes.FailedPrecondition, "tidb cluster is not available")
		}

		tidbCluster.SharedPoolId = currTiDBCluster.GetSharedPoolId()
		tidbCluster.TenantId = currTiDBCluster.GetTenantId()
		tidbCluster.Phase = currTiDBCluster.GetPhase()
		tidbCluster.Version = currTiDBCluster.GetVersion()

		// check if the sharedPool is locked
		unlockFunc, err := s.lock(ctx, currTiDBCluster.GetSharedPoolId())
		if err != nil {
			tilog.Error(ctx, "failed to lock sharedPool", zap.String("sharedPoolId", currTiDBCluster.GetSharedPoolId()), zap.Error(err))
			return err
		}
		defer func() {
			if err := unlockFunc(); err != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", currTiDBCluster.GetSharedPoolId()), zap.Error(err))
			}
		}()

		// get shared first
		sharedPool, err := s.sharedPoolRepo.Load(ctx, tidbCluster.GetSharedPoolId(), false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
			}
			// not found, check if it's in creating
			return status.Errorf(codes.NotFound, "sharedPool %s not found", tidbCluster.GetSharedPoolId())
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return err
		}
		// merge sharedPool status from dag task
		err = s.mergeSharedPoolStatusFromDagTask(ctx, sharedPool)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to merge sharedPool status from dag task: %s", err)
		}

		// get running tidb cluster from dag
		dagIns, err := s.getRunningDagIns(ctx, []string{
			dag.PlanCreateTiDBCluster,
			dag.PlanUpdateTiDBCluster,
			dag.PlanDeleteTiDBCluster,
		}, []string{tidbCluster.GetSharedPoolId()}, nil, []string{tidbCluster.GetId()}, nil)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get running tidb cluster dag instance: %s", err)
		}

		for _, dagIns := range dagIns {
			switch dagIns.ExecTask.RootName {
			case dag.PlanCreateTiDBCluster:
				return status.Errorf(codes.FailedPrecondition, "tidb cluster is creating, please try again later")
			case dag.PlanUpdateTiDBCluster:
				return status.Errorf(codes.FailedPrecondition, "tidb cluster is updating, please try again later")
			case dag.PlanDeleteTiDBCluster:
				return status.Errorf(codes.FailedPrecondition, "tidb cluster is deleting")
			}
		}

		// check request unit
		if tidbCluster.GetRequestUnit().GetMax() < tidbCluster.GetRequestUnit().GetMin() {
			return status.Errorf(codes.InvalidArgument, "request unit min %d cannot be greater than max %d", tidbCluster.GetRequestUnit().GetMin(), tidbCluster.GetRequestUnit().GetMax())
		}

		ruConfig, err := s.ruConfigRepo.GetRuConfig(ctx, sharedPool.GetId(), apiv1.RuConfig_SHARED_POOL)
		if err != nil {
			tilog.Error(ctx, "failed to get ru config", zap.String("sharedPoolId", sharedPool.GetId()), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to get ru config: %s", err)
		}

		optionsMap := map[string][]string{}
		componentTypes := []string{}
		componentIds := []string{}
		oldRequestUnit := &entity.RequestUnit{
			Min: currTiDBCluster.GetRequestUnit().GetMin(),
			Max: currTiDBCluster.GetRequestUnit().GetMax(),
		}
		oldRequestUnitStr := requestUnitToString(currTiDBCluster.GetRequestUnit())
		newRequestUnitStr := requestUnitProtoToString(tidbCluster.GetRequestUnit())
		essentialV2 := entity.IsEssentialV2(sharedPool)
		configuredScaleSharedPoolByRu := ruConfig.GetFeatureFlags().GetEnableScaleSharedPoolByRu()
		configuredScaleTiDBGroupByRu := ruConfig.GetFeatureFlags().GetEnableScaleTidbGroupByRu()
		enableScaleSharedPoolByRu := configuredScaleSharedPoolByRu && !essentialV2
		enableScaleTiDBGroupByRu := configuredScaleTiDBGroupByRu && !essentialV2
		tilog.Info(ctx, "resolved update tidb cluster RU spec scaling policy",
			zap.String("tidbClusterId", tidbCluster.GetId()),
			zap.String("sharedPoolId", sharedPool.GetId()),
			zap.Bool("essentialV2", essentialV2),
			zap.Bool("configuredEnableScaleSharedPoolByRu", configuredScaleSharedPoolByRu),
			zap.Bool("effectiveEnableScaleSharedPoolByRu", enableScaleSharedPoolByRu),
			zap.Bool("configuredEnableScaleTiDBGroupByRu", configuredScaleTiDBGroupByRu),
			zap.Bool("effectiveEnableScaleTiDBGroupByRu", enableScaleTiDBGroupByRu),
		)

		var nodeSpecMap map[string][]*entity.NodeSpec
		if !essentialV2 {
			nodeSpecMap, err = s.genNodeSpecMap(ctx, sharedPool)
			if err != nil {
				tilog.Error(ctx, "failed to get node spec map", zap.String("dataplaneId", sharedPool.GetDataplaneName()), zap.Error(err))
				return status.Errorf(codes.Internal, "failed to get node spec map: %s", err)
			}
		}

		// if set higher ru, we need re generate shared pool node spec
		if enableScaleSharedPoolByRu &&
			(currTiDBCluster.GetRequestUnit().GetMax() < tidbCluster.GetRequestUnit().GetMax() ||
				currTiDBCluster.GetRequestUnit().GetMin() < tidbCluster.GetRequestUnit().GetMin()) {

			ruTotal, err := s.getRequestUnitTotal(ctx, tidbCluster)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to get request unit total: %s", err)
			}

			// Snapshot old specs for whitelist components before regeneration,
			// so ScaleOverride can compare old vs new.
			oldSpecs := snapshotComponentSpecs(sharedPool)

			// force must not bypass the tikv-remote-worker topology guard - it only
			// relaxes the AVAILABLE-phase check inside generatedSharedPoolComponent.
			if err := s.guardTiKVRemoteWorkersTopologyStable(ctx, sharedPool); err != nil {
				return err
			}

			err = s.generatedSharedPoolComponent(ctx, sharedPool, ruTotal, nodeSpecMap, ruConfig, nil, force)
			if err != nil {
				return err
			}

			tilog.Info(ctx, "update tidb cluster request unit, re-generate shared pool node spec",
				zap.String("oldRequestUnit", oldRequestUnitStr),
				zap.String("newRequestUnit", newRequestUnitStr),
				zap.Any("ruTotal", ruTotal),
				zap.String("sharedPoolId", sharedPool.GetId()),
				zap.Bool("configuredEnableScaleSharedPoolByRu", configuredScaleSharedPoolByRu),
				zap.Bool("effectiveEnableScaleSharedPoolByRu", enableScaleSharedPoolByRu),
			)

			rcuSet := rcuScaleComponentSet()

			// Only submit actions for whitelist SharedPool components.
			for _, comp := range sharedPool.GetComponentGroups() {
				cfg, inWhitelist := rcuSet[comp.GetComponentType()]
				if !inWhitelist || cfg.Scope != RcuScaleScopeSharedPool {
					continue
				}
				if comp.GetPhase() != apiv1.Component_MODIFYING {
					continue
				}
				// If a ScaleOverride is configured, check whether the
				// old→new transition is allowed.
				if cfg.ScaleOverride != nil {
					key := componentSpecKey{compType: comp.GetComponentType(), id: comp.GetId()}
					oldSpec, found := oldSpecs[key]
					if found && !cfg.ScaleOverride(oldSpec, comp.GetNodeQuota().GetNodeSpec()) {
						tilog.Info(ctx, "rcu_scale_up: override vetoed scale-up",
							zap.String("sharedPoolId", sharedPool.GetId()),
							zap.String("tidbClusterId", tidbCluster.GetId()),
							zap.String("componentType", comp.GetComponentType().String()),
							zap.String("componentId", comp.GetId()))
						AutoscaleRcuScaleUpTotal.WithLabelValues(sharedPool.GetId(), "update", comp.GetComponentType().String(), "vetoed").Inc()
						continue
					}
				}

				optionsMap[entity.GetUpdateComponentKey(comp)] = []string{
					string(tioperator.OpsUpdateNodeSpec),
					string(tioperator.OpsUpdateNodeCount),
				}

				componentTypes = append(componentTypes, entity.GetComponentTypeName(comp))
				componentIds = append(componentIds, comp.GetId())
				AutoscaleRcuScaleUpTotal.WithLabelValues(sharedPool.GetId(), "update", comp.GetComponentType().String(), "prepared").Inc()
			}

		} else {
			tilog.Info(ctx, "update tidb cluster request unit, but not re-generate shared pool node spec",
				zap.String("oldRequestUnit", oldRequestUnitStr),
				zap.String("newRequestUnit", newRequestUnitStr),
				zap.String("sharedPoolId", sharedPool.GetId()),
				zap.Bool("configuredEnableScaleSharedPoolByRu", configuredScaleSharedPoolByRu),
				zap.Bool("effectiveEnableScaleSharedPoolByRu", enableScaleSharedPoolByRu),
			)
		}

		currTiDBCluster.RequestUnit = tidbCluster.GetRequestUnit()

		var needSendMeteringMetadata = true

		// update ru_version if set (> 0)
		if tidbCluster.GetRuVersion() > 0 {
			if currTiDBCluster.GetRuVersion() != tidbCluster.GetRuVersion() {
				// ru_version changed, must send metering metadata so metering gets the latest ru_version
				needSendMeteringMetadata = true
			}
			currTiDBCluster.RuVersion = tidbCluster.GetRuVersion()
		}
		// Evaluate per-cluster component specs against the new RU.
		if enableScaleTiDBGroupByRu {
			tilog.Info(ctx, "update tidb cluster request unit, evaluate tidb group node spec",
				zap.String("tidbClusterId", tidbCluster.GetId()),
				zap.String("oldRequestUnit", oldRequestUnitStr),
				zap.String("newRequestUnit", newRequestUnitStr),
				zap.Bool("configuredEnableScaleTiDBGroupByRu", configuredScaleTiDBGroupByRu),
				zap.Bool("effectiveEnableScaleTiDBGroupByRu", enableScaleTiDBGroupByRu),
			)

			// list tidb groups
			tidbGroups, err := s.tidbGroupRepo.List(ctx, []string{sharedPool.GetId()}, nil, []string{currTiDBCluster.GetId()}, nil, false)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to list tidb groups: %s", err)
			}

			updatedTiDBComponentIDs := make(map[string]struct{}, len(tidbGroups))
			for _, tidbGroup := range tidbGroups {
				updated, err := s.regenerateTiDBGroupForRUScaleUp(
					ctx,
					sharedPool,
					oldRequestUnit,
					tidbCluster,
					tidbGroup,
					nodeSpecMap,
					ruConfig,
				)
				if err != nil {
					return err
				}
				if updated {
					updatedTiDBComponentIDs[tidbGroup.GetId()] = struct{}{}
				}
			}

			if !entity.HasSharedTiProxy(sharedPool) && entity.FindTiProxyComponent(sharedPool, tidbCluster.GetId()) != nil {
				err := s.generatedTiProxyGroup(ctx, sharedPool, tidbCluster, nodeSpecMap, ruConfig)
				if err != nil {
					return err
				}
			}

			// Submit only components changed by this request. This avoids
			// resubmitting an unrelated TiDB component that was already MODIFYING.
			for _, comp := range sharedPool.GetComponentGroups() {
				isTiDBComponent := comp.GetComponentType() == apiv1.ComponentType_TIDB
				isPerClusterTiProxyComponent := !entity.HasSharedTiProxy(sharedPool) &&
					comp.GetComponentType() == apiv1.ComponentType_TIPROXY
				if !isTiDBComponent && !isPerClusterTiProxyComponent {
					continue
				}

				if isTiDBComponent {
					if _, updated := updatedTiDBComponentIDs[comp.GetId()]; !updated {
						continue
					}
					optionsMap[entity.GetUpdateComponentKey(comp)] = []string{
						string(tioperator.OpsUpdateNodeSpec),
					}
				} else {
					if comp.GetPhase() != apiv1.Component_MODIFYING {
						continue
					}
					optionsMap[entity.GetUpdateComponentKey(comp)] = []string{
						string(tioperator.OpsUpdateNodeSpec),
						string(tioperator.OpsUpdateNodeCount),
					}
				}

				componentTypes = append(componentTypes, entity.GetComponentTypeName(comp))
				componentIds = append(componentIds, comp.GetId())
				AutoscaleRcuScaleUpTotal.WithLabelValues(
					sharedPool.GetId(),
					"update",
					comp.GetComponentType().String(),
					"prepared",
				).Inc()
			}

			needSendMeteringMetadata = true
			tidbCluster.Phase = apiv1.Cluster_MODIFYING
		} else {
			tilog.Info(ctx, "update tidb cluster request unit, but not re-generate tidb group node spec",
				zap.String("tidbClusterId", currTiDBCluster.GetId()),
				zap.String("oldRequestUnit", oldRequestUnitStr),
				zap.String("newRequestUnit", newRequestUnitStr),
				zap.Bool("configuredEnableScaleTiDBGroupByRu", configuredScaleTiDBGroupByRu),
				zap.Bool("effectiveEnableScaleTiDBGroupByRu", enableScaleTiDBGroupByRu),
			)
		}
		// save tidb cluster new ru
		err = s.tidbClusterRepo.Save(ctx, currTiDBCluster)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to save tidbCluster: %s", err)
		}

		sharedPoolStr, err := protojson.Marshal(sharedPool)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to marshal shard pool: %s", err)
		}

		args := actions.ActionArgs{
			workflow.SharedPoolIdKey:   actions.ActionArgString(tidbCluster.GetSharedPoolId()),
			workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
			workflow.ComponentIdsKey:   actions.ActionArgString(strings.Join(componentIds, ",")),
			workflow.ComponentTypesKey: actions.ActionArgString(strings.Join(componentTypes, ",")),
			workflow.TenantIdKey:       actions.ActionArgString(tidbCluster.GetTenantId()),
			workflow.TiDBClusterIdKey:  actions.ActionArgString(tidbCluster.GetId()),
			utils.TraceIdKey:           utils.GetTraceId(ctx),
		}
		if essentialV2 && len(componentIds) == 0 && len(componentTypes) == 0 {
			args[workflow.SkipComponentHPAApplyKey] = actions.ActionArgString(workflow.SkipComponentHPAApplyValue)
		}
		if needSendMeteringMetadata {
			args[workflow.MeteringTypeKey] = actions.ActionArgString(workflow.MeteringTypeValueLogicMetaByTiDBCluster)
		}

		// build update options
		for k, v := range optionsMap {
			args[k] = actions.ActionArgString(strings.Join(v, ","))
		}
		tilog.Info(ctx, "rcu_scale_up: prepared update tidb cluster autoscale output",
			zap.String("sharedPoolId", tidbCluster.GetSharedPoolId()),
			zap.String("tidbClusterId", tidbCluster.GetId()),
			zap.String("oldRequestUnit", oldRequestUnitStr),
			zap.String("newRequestUnit", newRequestUnitStr),
			zap.String("componentIds", strings.Join(componentIds, ",")),
			zap.String("componentTypes", strings.Join(componentTypes, ",")),
			zap.String("components", describeComponentSpecsByIDs(sharedPool, componentIds)),
			zap.String("options", describeOptionsMap(optionsMap)),
			zap.Bool("needSendMeteringMetadata", needSendMeteringMetadata),
		)
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("update-tidbcluster-%s-%s-%s",
			tidbCluster.GetSharedPoolId(),
			tidbCluster.GetId(),
			utils.FnvHash(string(sharedPoolStr)),
		))
		err = s.dagExecutor.Submit(ctx, dag.PlanUpdateTiDBCluster, args, dagExecId)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to submit update tidb cluster dag: %s", err)
		}
		dagSubmitted = true

		return nil
	}); err != nil {
		if dagSubmitted {
			if abortErr := s.dagExecutor.Abort(ctx, dagExecId); abortErr != nil {
				tilog.Error(ctx, "failed to abort dag after transaction failure",
					zap.String("dagExecId", dagExecId),
					zap.Error(abortErr))
			}
		}
		return nil, err
	}

	return tidbCluster, nil
}