package cluster

import (
	"context"
	"fmt"
	"slices"
	"strings"

	"github.com/samber/lo"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"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/entity"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/tidbgroup"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/gkeheterogeneous"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/workflow/action"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/workflow/common"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/workflow/dag"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/clusterevent"
)

func (s *serviceImpl) Scale(
	ctx context.Context,
	request *globaldedicatedv1.ScaleClusterWithMultiComponentRequest,
) (*entity.Cluster, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "Domain.Scale")
	defer span.End()

	cluster, err := s.clusterRepo.GetAndLoadComponents(ctx, request.ClusterId)
	if err != nil {
		return nil, err
	}

	regionEn, err := s.regionRepo.SelectRegionById(ctx, cluster.RegionId, cluster.OrganizationId)
	if err != nil {
		return nil, err
	}
	if regionEn.State != globalcommonv1.Region_ACTIVE {
		return nil, bizerr.New(bizerr.ErrCodePermission, "region is not active")
	}

	if cluster.State != globaldedicatedv1.Cluster_AVAILABLE && !request.ForceScale {
		return nil, bizerr.New(bizerr.ErrCodeParams, "cluster is not available for this action")
	}

	// Not safe, could have concurrency issue if two scale requests come in at the same time.
	// FIXME: Should be implemented with db row lock.
	if ok := dag.CanRunScaleDag(ctx, s.workflow, cluster.ClusterId); !ok {
		return nil, bizerr.New(bizerr.ErrCodeParams, "cluster is not available for this action")
	}

	tidbGroups, err := s.tidbGroupRepo.Query(ctx, &tidbgroup.Filter{
		ClusterIds: []uint64{cluster.ClusterId},
	})
	if err != nil {
		return nil, err
	}
	if len(tidbGroups) == 0 {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "empty tidb group")
	}

	if len(request.ClusterComponents) == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "empty cluster components")
	}

	// Check whether the cloud provider is supported for the cluster.
	if err := checkSupportedCloudProvider(request.GetClusterComponents(), cluster.GetCloudProviderId()); err != nil {
		return nil, err
	}

	cluster.DesiredComponents, err = s.buildDesiredComponents(
		ctx,
		tidbGroups,
		cluster,
		request.ClusterComponents,
		false,
	)
	if err != nil {
		return nil, err
	}

	if err := s.CheckCreateDesiredComponents(ctx, cluster, cluster.DesiredComponents); err != nil {
		return nil, err
	}

	// Always check TiProxy replica count for scale operations (existing clusters)
	if err := tiproxyReplicaCountCheck(cluster); err != nil {
		return nil, err
	}

	if err := s.CheckScaleDesiredComponents(ctx, cluster, tidbGroups, cluster.DesiredComponents, regionEn); err != nil {
		return nil, err
	}

	for _, desired := range cluster.DesiredComponents {
		if cluster.GetCloudProviderId() != globalcommonv1.CloudProvider_GCP {
			continue
		}
		if !gkeheterogeneous.IsSupportedRegionalHeterogeneousComponentType(desired.GetComponentType()) {
			continue
		}
		currentComp, _ := lo.Find(cluster.GetClusterComponents(), func(c *globaldedicatedv1.Cluster_ClusterComponent) bool {
			return c.GetComponentType() == desired.GetComponentType() && c.GetTidbGroupId() == desired.GetTidbGroupId()
		})
		if mu := instanceMuChanged(ctx, cluster, currentComp, desired); mu != nil && mu.InstanceMuChanged {
			if berr := s.checkNoApplyHeterogeneousDeployPending(
				ctx, cluster.GetClusterId(), desired.GetComponentType(), desired.GetTidbGroupId(),
			); berr != nil {
				return nil, berr
			}
		}
	}

	cluster.State = globaldedicatedv1.Cluster_SCALING
	cluster.HealthState = globaldedicatedv1.Cluster_HEALTHY

	// validate only
	if request.ValidateOnly {
		mergeClusterComponent(cluster, cluster.DesiredComponents)
		return cluster, nil
	}

	if len(cluster.DesiredComponents) == 0 {
		tilog.Info(ctx, "no component needs to perform scale operations")
		return cluster, nil
	}

	defaultTidbGroupId, err := getDefaultTidbGroup(ctx, s.tidbGroupRepo, cluster.ClusterId)
	if err != nil {
		return nil, err.Chain("query default tidb node group failed: cluster=%d, err=%w", cluster.ClusterId, err)
	}

	operationContext := entity.BuildScaleOperationContext(
		request.GetSourceOperation(),
		request.GetRequestedFields(),
		request.GetClusterComponents(),
	)
	var newCluster *entity.Cluster
	err_ := s.transaction.Transaction(ctx, func(ctx context.Context) error {
		// save desired components, and build scale workflow tasks.
		scaleDagTasks, err := s.saveDesiredComponents(ctx, cluster, regionEn, cluster.DesiredComponents, defaultTidbGroupId)
		if err != nil {
			return err.Chain("apply desired components failed: desired=%+v, err=%w", cluster.DesiredComponents, err)
		}
		scaleDagTasks.OperationContext = operationContext

		// update cluster state
		if err := s.clusterRepo.Save(ctx, cluster); err != nil {
			return err
		}

		newCluster, err = s.clusterRepo.GetAndLoadComponents(ctx, cluster.ClusterId)
		if err != nil {
			return err.Chain("load cluster failed: cluster=%d, err=%w", cluster.ClusterId, err)
		}
		beforeTopology := entity.BuildClusterTopology(cluster, defaultTidbGroupId)
		afterTopology := entity.BuildClusterTopology(newCluster, defaultTidbGroupId)
		scaleEvent := entity.NewClusterEventBuilder(ctx, cluster, clusterevent.EventTypeScaling).
			SetTopologyBefore(beforeTopology).
			SetTopologyAfter(afterTopology).
			SetOperationContext(operationContext).
			Build()
		if err := s.eventCommitter.Write(ctx, clusterevent.TopicClusterEvent, scaleEvent); err != nil {
			return err
		}

		// run workflow
		// the request allow scale multiple components at the same time. However, we only support scaling one component at a time to ensure the safety and simplicity of the cluster.
		// This transpiling algorithm is implemented in the [dag.BuildScaleTopology] method.
		if err := dag.RunScaleDag(ctx, s.workflow, scaleDagTasks, request.ForceScale); err != nil {
			return bizerr.WrapError(err, "run dag failed")
		}

		return nil
	})
	if err_ != nil {
		return nil, bizerr.WrapGormError(err_, "scale cluster failed")
	}

	return newCluster, nil
}

func mergeClusterComponent(cluster *entity.Cluster, desiredComponents []*entity.DesiredClusterComponent) {
nextDesiredComp:
	for _, desire := range desiredComponents {
		for i, current := range cluster.GetClusterComponents() {
			// Replace the desired components.
			if desire.GetComponentType() == current.GetComponentType() && desire.GetTidbGroupId() == current.GetTidbGroupId() {
				cluster.GetClusterComponents()[i] = desire.ClusterComponent
				continue nextDesiredComp
			}
		}
		// If not found (not continue on the previous line), add the new components.
		if desire.GetId() == 0 {
			cluster.ClusterComponents = append(cluster.ClusterComponents, desire.ClusterComponent)
		}
	}
}

func getDefaultTidbGroup(
	ctx context.Context,
	tidbGroupRepo tidbgroup.Repo,
	clusterId uint64,
) (uint64, bizerr.BizError) {
	wp, err := tidbGroupRepo.LoadDefault(ctx, clusterId)
	if err != nil {
		return 0, err
	}

	return wp.Id, nil
}

func (s *serviceImpl) saveDesiredComponents(
	ctx context.Context,
	cluster *entity.Cluster,
	regionEn *entity.Region,
	desiredComponents []*entity.DesiredClusterComponent,
	defaultTidbGroupId uint64,
) (*dag.ScaleDagTasks, bizerr.BizError) {
	out := &dag.ScaleDagTasks{
		TenantId:                         cluster.GetOrganizationId(),
		ProjectId:                        cluster.GetProjectId(),
		ClusterId:                        cluster.GetClusterId(),
		CloudProviderId:                  cluster.GetCloudProviderId(),
		ComponentsAsConfigOrigin:         make([]*action.ScaleComponentParameter, 0),
		ComponentsAsConfigFollower:       make([]*action.ScaleComponentParameter, 0),
		ComponentTiProxyAsConfigOrigin:   make([]*action.ScaleComponentParameter, 0),
		ComponentTiProxyAsConfigFollower: make([]*action.ScaleComponentParameter, 0),
	}

	// Analyze scale operations
	out.ScaleOperations = s.analyzeScaleOperations(ctx, cluster, desiredComponents)

	grouped := lo.GroupBy(desiredComponents, func(i *entity.DesiredClusterComponent) globaldedicatedv1.ComponentType {
		return i.GetComponentType()
	})

	// calculate scaling options
	var options []string
	for k, v := range s.cfg.BizConfig.ClusterService.ScaleOptions {
		if v.CheckWithAll(cluster.OrganizationId, cluster.ProjectId, cluster.CloudProviderId.String(), cluster.RegionId.String()) {
			options = append(options, k)
		}
	}
	for _, components := range grouped {
		slices.SortFunc(components, func(a, b *entity.DesiredClusterComponent) int {
			if a.GetTidbGroupId() == defaultTidbGroupId {
				return -1
			}

			if b.GetTidbGroupId() == defaultTidbGroupId {
				return 1
			}

			if a.Id == 0 {
				return 1
			}

			if b.Id == 0 {
				return -1
			}

			if a.Id < b.Id {
				return -1
			}
			if a.Id > b.Id {
				return 1
			}
			return 0
		})

		for i, comp := range components {
			currentComp, _ := lo.Find(cluster.GetClusterComponents(), func(c *entity.ClusterComponent) bool {
				return c.GetComponentType() == comp.GetComponentType() && c.GetTidbGroupId() == comp.GetTidbGroupId()
			})

			// save storage scale operation
			if err := s.saveStorageOperations(ctx, cluster.ClusterId, currentComp, comp); err != nil {
				return nil, err
			}

			// save components
			id, err := s.componentRepo.Save(ctx, comp)
			if err != nil {
				return nil, err
			}

			task := &action.ScaleComponentParameter{
				ID:                    id,
				ComponentType:         comp.GetComponentType().String(),
				NeedToScaleInstanceMu: instanceMuChanged(ctx, cluster, currentComp, comp),
				Options:               options,
			}

			if task.NeedToScaleInstanceMu != nil && task.NeedToScaleInstanceMu.InstanceMuChanged &&
				cluster.GetCloudProviderId() == globalcommonv1.CloudProvider_GCP {
				componentType := common.ScaleComponentTypeFromString(task.ComponentType)
				if gkeheterogeneous.IsSupportedRegionalHeterogeneousComponentType(componentType) {
					regionalCompType, _ := gkeheterogeneous.ToRegionalComponentType(componentType)
					tidbGroupID := ""
					if comp.GetTidbGroupId() != 0 {
						tidbGroupID = fmt.Sprintf("%d", comp.GetTidbGroupId())
					}
					_, deployExists, berr := gkeheterogeneous.GetDeploy(
						ctx, s.regionServerCli, regionEn,
						fmt.Sprintf("%d", cluster.GetClusterId()),
						regionalCompType, tidbGroupID,
					)
					if berr != nil {
						return nil, berr
					}
					if deployExists {
						task.Options = append(append([]string{}, options...), dag.OptionDisableGKEHeterogeneous)
					}
				}
			}

			if i == 0 && task.NeedToScaleInstanceMu != nil && task.NeedToScaleInstanceMu.InstanceMuChanged {
				if comp.GetComponentType() == globaldedicatedv1.ComponentType_TIPROXY {
					out.ComponentTiProxyAsConfigOrigin = append(out.ComponentTiProxyAsConfigOrigin, task)
				} else {
					out.ComponentsAsConfigOrigin = append(out.ComponentsAsConfigOrigin, task)
				}
			} else {
				if comp.GetComponentType() == globaldedicatedv1.ComponentType_TIPROXY {
					out.ComponentTiProxyAsConfigFollower = append(out.ComponentTiProxyAsConfigFollower, task)
				} else {
					out.ComponentsAsConfigFollower = append(out.ComponentsAsConfigFollower, task)
				}
			}
		}
	}

	return out, nil
}

func (s *serviceImpl) saveStorageOperations(
	ctx context.Context,
	clusterId uint64,
	currentComp *globaldedicatedv1.Cluster_ClusterComponent,
	desiredComp *entity.DesiredClusterComponent,
) bizerr.BizError {
	for diskType := range storageChanged(currentComp, desiredComp) {
		err := s.operationRepo.Save(ctx, &entity.ClusterOperation{
			ClusterID: clusterId,
			Type:      globaldedicatedv1.ClusterOperationType_SCALE,
			Operation: &entity.ScaleOperation{
				ComponentID:    desiredComp.Id,
				StorageChanged: true,
				ComponentType:  currentComp.GetComponentType(),
				StorageType:    diskType,
			},
			State: globaldedicatedv1.ClusterOperationState_CREATED,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func instanceMuChanged(
	ctx context.Context,
	cluster *entity.Cluster,
	currentComp *entity.ClusterComponent,
	desiredComp *entity.DesiredClusterComponent,
) *action.NeedToScaleInstanceMu {
	if currentComp == nil {
		compFromOtherGroups := lo.Filter(cluster.GetClusterComponents(), func(c *entity.ClusterComponent, _ int) bool {
			return c.GetComponentType() == desiredComp.GetComponentType()
		})
		instanceChanged := true
		if len(compFromOtherGroups) > 0 {
			instanceChanged = compFromOtherGroups[0].GetSkuKey() != desiredComp.Sku.GetKey()
		}

		return &action.NeedToScaleInstanceMu{
			InstanceMuChanged: instanceChanged,
			ScaleUp:           true,
			ScaleUpFromZero:   true,
		}
	}

	if currentComp.GetCount() == 0 {
		// todo: add metrics
		tilog.Error(
			ctx,
			"unexpect component value in db",
			zap.Uint64("id", currentComp.Id),
			zap.String("component", currentComp.GetComponentType().String()),
		)
		return &action.NeedToScaleInstanceMu{
			InstanceMuChanged: true,
			ScaleUp:           true,
			ScaleUpFromZero:   true,
		}
	}

	if currentComp.GetSkuKey() == desiredComp.Sku.GetKey() {
		return nil
	}

	if desiredComp.Sku.GetInstanceType() == currentComp.GetInstanceType() {
		return nil
	}

	if desiredComp.Sku.GetCpuCore() > int32(currentComp.GetCpuNum()) {
		return &action.NeedToScaleInstanceMu{
			InstanceMuChanged: true,
			ScaleUp:           true,
			CurrentInstanceMu: currentComp.GetInstanceType(),
			TargetInstanceMu:  desiredComp.Sku.GetInstanceType(),
		}
	}

	if desiredComp.Sku.GetCpuCore() < int32(currentComp.GetCpuNum()) {
		return &action.NeedToScaleInstanceMu{
			InstanceMuChanged: true,
			ScaleUp:           false,
			CurrentInstanceMu: currentComp.GetInstanceType(),
			TargetInstanceMu:  desiredComp.Sku.GetInstanceType(),
		}
	}

	if desiredComp.Sku.GetMemoryGi() > int32(currentComp.GetMemorySizeGi()) {
		return &action.NeedToScaleInstanceMu{
			InstanceMuChanged: true,
			ScaleUp:           true,
			CurrentInstanceMu: currentComp.GetInstanceType(),
			TargetInstanceMu:  desiredComp.Sku.GetInstanceType(),
		}
	}

	return &action.NeedToScaleInstanceMu{
		InstanceMuChanged: true,
		ScaleUp:           false,
		CurrentInstanceMu: currentComp.GetInstanceType(),
		TargetInstanceMu:  desiredComp.Sku.GetInstanceType(),
	}
}

func storageChanged(
	currentComp *globaldedicatedv1.Cluster_ClusterComponent,
	desiredComp *entity.DesiredClusterComponent,
) map[globaldedicatedv1.Sku_DiskType]bool {
	out := make(map[globaldedicatedv1.Sku_DiskType]bool)

	if _, ok := entity.ComponentsHaveStorage[desiredComp.GetComponentType()]; !ok {
		return out
	}

	customizedStorageParameterChanged := func(currentComp *entity.ClusterComponent, desiredComp *entity.DesiredClusterComponent, diskType globaldedicatedv1.Sku_DiskType) bool {
		return currentComp.GetStorages()[diskType.String()].GetSizeGi() != desiredComp.GetStorages()[diskType.String()].GetSizeGi() ||
			currentComp.GetStorages()[diskType.String()].GetIops() != desiredComp.GetStorages()[diskType.String()].GetIops() ||
			currentComp.GetStorages()[diskType.String()].GetThroughputMibps() != desiredComp.GetStorages()[diskType.String()].GetThroughputMibps() ||
			currentComp.GetStorages()[diskType.String()].GetVolumeType() != desiredComp.GetStorages()[diskType.String()].GetVolumeType()
	}

	if customizedStorageParameterChanged(currentComp, desiredComp, globaldedicatedv1.Sku_DATA) {
		out[globaldedicatedv1.Sku_DATA] = true
	}

	if customizedStorageParameterChanged(currentComp, desiredComp, globaldedicatedv1.Sku_RAFT_LOG) {
		out[globaldedicatedv1.Sku_RAFT_LOG] = true
	}

	return out
}

// analyzeScaleOperation analyzes scale operation type for each component type
func (s *serviceImpl) analyzeScaleOperations(
	ctx context.Context,
	cluster *entity.Cluster,
	desiredComponents []*entity.DesiredClusterComponent,
) map[string][]string {
	operations := make(map[string][]string)
	if len(desiredComponents) == 0 {
		return operations
	}

	// Record operation type for each component type
	for _, desiredComp := range desiredComponents {
		// Find current component
		currentComp, found := lo.Find(cluster.GetClusterComponents(), func(c *entity.ClusterComponent) bool {
			return c.GetComponentType() == desiredComp.GetComponentType() && c.GetTidbGroupId() == desiredComp.GetTidbGroupId()
		})

		componentType := strings.ToLower(desiredComp.GetComponentType().String())
		if !found {
			// New component - scale out
			operations[componentType] = append(operations[componentType], "node_scale_out")
		} else {
			// Compare component specification changes
			switch {
			case currentComp.GetCount() == 0 && desiredComp.GetCount() > 0:
				// Scale out from 0 - scale out
				operations[componentType] = append(operations[componentType], "node_scale_out")
			case currentComp.GetCount() > 0 && desiredComp.GetCount() == 0:
				// Scale in to 0 - scale in
				operations[componentType] = append(operations[componentType], "node_scale_in")
			case currentComp.GetSkuKey() != desiredComp.Sku.GetKey():
				// Specification changes
				if desiredComp.Sku.GetCpuCore() > int32(currentComp.GetCpuNum()) ||
					desiredComp.Sku.GetMemoryGi() > int32(currentComp.GetMemorySizeGi()) {
					// Specification upgrade - scale up
					operations[componentType] = append(operations[componentType], "node_scale_up")
				} else {
					// Specification downgrade - scale down
					operations[componentType] = append(operations[componentType], "node_scale_down")
				}
			case currentComp.GetCount() != desiredComp.GetCount():
				// Node count changes
				if desiredComp.GetCount() > currentComp.GetCount() {
					operations[componentType] = append(operations[componentType], "node_scale_out")
				} else {
					operations[componentType] = append(operations[componentType], "node_scale_in")
				}
			}
			// Check for disk/storage changes
			diskChanged := storageChanged(currentComp, desiredComp)
			if len(diskChanged) > 0 {
				// Check if disk size/iops is increasing or decreasing
				for diskType := range diskChanged {
					currentStorages := currentComp.GetStorages()
					desiredStorages := desiredComp.GetStorages()
					if currentStorages == nil || desiredStorages == nil {
						continue
					}
					currentStorage, currentExists := currentStorages[diskType.String()]
					desiredStorage, desiredExists := desiredStorages[diskType.String()]
					if !currentExists || !desiredExists {
						continue
					}
					if desiredStorage.GetSizeGi() > currentStorage.GetSizeGi() {
						operations[componentType] = append(operations[componentType], "disk_size_scale_up")
					} else if desiredStorage.GetSizeGi() < currentStorage.GetSizeGi() {
						operations[componentType] = append(operations[componentType], "disk_size_scale_down")
					}
					if desiredStorage.GetIops() > currentStorage.GetIops() {
						operations[componentType] = append(operations[componentType], "disk_iops_scale_up")
					} else if desiredStorage.GetIops() < currentStorage.GetIops() {
						operations[componentType] = append(operations[componentType], "disk_iops_scale_down")
					}
				}
			}
		}
	}
	for k := range operations {
		operations[k] = lo.Uniq(operations[k])
	}

	tilog.Info(ctx, "analyzed scale operation", zap.Any("component_operations", operations))
	return operations
}