package server

import (
	"context"
	"slices"
	"sort"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	scaleConstraintRuleTypeNone       = "NONE"
	scaleConstraintRuleTypeMultipleOf = "MULTIPLE_OF"
	scaleConstraintRuleTypeFixed      = "FIXED"

	scaleConstraintReasonComponentPhaseUnavailable = "COMPONENT_PHASE_NOT_AVAILABLE"
	scaleConstraintReasonBelowRURequirement        = "BELOW_RU_REQUIREMENT"
	scaleConstraintReasonInsufficientLocalDisk     = "INSUFFICIENT_LOCAL_DISK_SIZE"
	scaleConstraintReasonSpecNotPresent            = "SPEC_NOT_PRESENT_IN_SHARED_POOL"
)

var scaleConstraintComponentTypes = map[apiv1.ComponentType]struct{}{
	apiv1.ComponentType_TIDB:               {},
	apiv1.ComponentType_TIKV:               {},
	apiv1.ComponentType_PD:                 {},
	apiv1.ComponentType_PD_TSO:             {},
	apiv1.ComponentType_PD_SCHEDULING:      {},
	apiv1.ComponentType_TIPROXY:            {},
	apiv1.ComponentType_TIFLASH_WRITE:      {},
	apiv1.ComponentType_TIFLASH_COMPUTE:    {},
	apiv1.ComponentType_TIDB_WORKER:        {},
	apiv1.ComponentType_TIKV_WORKER:        {},
	apiv1.ComponentType_COPROCESSOR_WORKER: {},
	apiv1.ComponentType_RESOURCE_MANAGER:   {},
}

func (s *serverImpl) GetSharedPoolScaleConstraints(ctx context.Context, sharedPoolId string) (*entity.GetSharedPoolScaleConstraintsResponse, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "sharedPoolId is empty")
	}

	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		if status.Code(err) == codes.NotFound {
			return nil, status.Errorf(codes.NotFound, "sharedPool %s not found", sharedPoolId)
		}
		return nil, status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
	}

	if err := s.mergeSharedPoolStatusFromDagTask(ctx, sharedPool); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to merge sharedPool status from dag task: %s", err)
	}

	nodeSpecMap, err := s.genNodeSpecMap(ctx, sharedPool)
	if err != nil {
		return nil, err
	}

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

	resp := &entity.GetSharedPoolScaleConstraintsResponse{
		SharedPoolId: sharedPool.GetId(),
		Components:   make([]*apiv1.SharedPoolScaleConstraintComponent, 0),
	}

	for _, comp := range sharedPool.GetComponentGroups() {
		if !isScaleConstraintComponentType(comp.GetComponentType()) {
			continue
		}

		minRU, skip, err := s.getMinRUForScaleConstraintComponent(ctx, sharedPool, comp)
		if err != nil {
			return nil, err
		}
		if skip {
			continue
		}

		if comp.GetNodeQuota() == nil || comp.GetNodeQuota().GetNodeSpec() == nil {
			return nil, status.Errorf(codes.Internal, "component %s has empty node quota", entity.GetComponentGroupName(comp))
		}

		out, err := s.buildSharedPoolScaleConstraintComponent(sharedPool, comp, nodeSpecMap, ruConfig, minRU)
		if err != nil {
			return nil, err
		}
		resp.Components = append(resp.Components, out)
	}

	return resp, nil
}

func isScaleConstraintComponentType(componentType apiv1.ComponentType) bool {
	_, ok := scaleConstraintComponentTypes[componentType]
	return ok
}

func (s *serverImpl) getMinRUForScaleConstraintComponent(
	ctx context.Context,
	sharedPool *entity.SharedPool,
	comp *entity.Component,
) (uint32, bool, error) {
	if comp.GetComponentType() != apiv1.ComponentType_TIDB {
		minRU, err := s.getMinRUForComponent(ctx, comp, sharedPool)
		return minRU, false, err
	}

	tidbGroup, err := s.tidbGroupRepo.Load(ctx, comp.GetId(), true)
	if err != nil {
		return 0, false, status.Errorf(codes.Internal, "failed to load tidb group %s including deleted records: %s", comp.GetId(), err)
	}
	if tidbGroup == nil {
		return 0, false, status.Errorf(codes.Internal, "loaded nil tidb group %s including deleted records", comp.GetId())
	}
	if tidbGroup.IsDeleted {
		return 0, true, nil
	}

	tidbCluster, err := s.tidbClusterRepo.Load(ctx, tidbGroup.GetClusterId(), false)
	if err != nil {
		return 0, false, status.Errorf(codes.Internal, "failed to load tidb cluster %s: %s", tidbGroup.GetClusterId(), err)
	}
	return tidbCluster.GetRequestUnit().GetMin(), false, nil
}

func (s *serverImpl) buildSharedPoolScaleConstraintComponent(
	sharedPool *entity.SharedPool,
	comp *entity.Component,
	nodeSpecMap map[string][]*entity.NodeSpec,
	ruConfig *entity.RuConfig,
	minRU uint32,
) (*apiv1.SharedPoolScaleConstraintComponent, error) {
	requiredMinSpec, hasRequiredMinSpec, err := getRequiredMinSpecForComponent(comp.GetComponentType(), ruConfig, minRU)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get min node spec from ru config: %s", err)
	}
	if !hasRequiredMinSpec {
		requiredMinSpec = nil
	}

	return &apiv1.SharedPoolScaleConstraintComponent{
		ComponentId:          comp.GetId(),
		ComponentType:        comp.GetComponentType(),
		BlockingReasons:      buildComponentBlockingReasons(sharedPool, comp),
		NodeCountConstraints: buildNodeCountConstraints(sharedPool, comp),
		NodeSpecConstraints:  buildNodeSpecConstraints(comp, nodeSpecMap, requiredMinSpec),
	}, nil
}

func buildComponentBlockingReasons(sharedPool *entity.SharedPool, comp *entity.Component) []*apiv1.SharedPoolScaleConstraintReason {
	if comp.GetPhase() == apiv1.Component_AVAILABLE {
		return nil
	}
	if canSkipPhaseBlockingForScaleConstraint(sharedPool, comp) {
		return nil
	}

	return []*apiv1.SharedPoolScaleConstraintReason{
		{
			Code:    scaleConstraintReasonComponentPhaseUnavailable,
			Message: "target component is not available",
		},
	}
}

func canSkipPhaseBlockingForScaleConstraint(sharedPool *entity.SharedPool, comp *entity.Component) bool {
	if comp == nil || comp.GetNodeQuota() == nil {
		return false
	}

	// Keep the query aligned with ScaleSharedPool's special-case phase bypass:
	// topology-driven workers only recover from an unavailable phase when the
	// current operation is effectively a cross-zero recovery from node_count=0.
	return isTopologyDrivenWorkerZeroCountAllowed(sharedPool, comp.GetComponentType()) &&
		comp.GetNodeQuota().GetNodeCount() == 0
}

func buildNodeCountConstraints(sharedPool *entity.SharedPool, comp *entity.Component) *apiv1.SharedPoolScaleConstraintNodeCountConstraints {
	current := comp.GetNodeQuota().GetNodeCount()
	minCount := comp.GetNodeQuota().GetMinNodeCount()
	maxCount := comp.GetNodeQuota().GetMaxNodeCount()
	allowCrossZero := false

	switch comp.GetComponentType() {
	case apiv1.ComponentType_TIFLASH_WRITE, apiv1.ComponentType_TIFLASH_COMPUTE:
		minCount = 0
		allowCrossZero = true
	case apiv1.ComponentType_TIDB_WORKER, apiv1.ComponentType_TIKV_WORKER:
		if isTopologyDrivenWorkerZeroCountAllowed(sharedPool, comp.GetComponentType()) {
			minCount = 0
			allowCrossZero = true
		}
	}

	return &apiv1.SharedPoolScaleConstraintNodeCountConstraints{
		Current:        current,
		Min:            minCount,
		Max:            maxCount,
		Rule:           buildNodeCountRule(comp.GetComponentType()),
		AllowCrossZero: allowCrossZero,
	}
}

func buildNodeCountRule(componentType apiv1.ComponentType) *apiv1.SharedPoolScaleConstraintNodeCountRule {
	switch componentType {
	case apiv1.ComponentType_TIKV, apiv1.ComponentType_PD:
		return &apiv1.SharedPoolScaleConstraintNodeCountRule{
			Type:  scaleConstraintRuleTypeMultipleOf,
			Value: 3,
		}
	case apiv1.ComponentType_PD_TSO:
		return &apiv1.SharedPoolScaleConstraintNodeCountRule{
			Type:  scaleConstraintRuleTypeFixed,
			Value: 2,
		}
	default:
		return &apiv1.SharedPoolScaleConstraintNodeCountRule{
			Type: scaleConstraintRuleTypeNone,
		}
	}
}

func buildNodeSpecConstraints(
	comp *entity.Component,
	nodeSpecMap map[string][]*entity.NodeSpec,
	requiredMinSpec *entity.NodeSpec,
) *apiv1.SharedPoolScaleConstraintNodeSpecConstraints {
	currentSpec := comp.GetNodeQuota().GetNodeSpec()
	currentName := entity.GetComponentNodeSpec(currentSpec)
	withLocalDisk := currentSpec.GetWithLocalDisk()
	storages := comp.GetStorages()

	keys := make([]string, 0, len(nodeSpecMap))
	hasCurrent := false
	for key, specs := range nodeSpecMap {
		if hasNodeSpecVariant(specs, withLocalDisk) {
			keys = append(keys, key)
			if key == currentName {
				hasCurrent = true
			}
		}
	}
	if !hasCurrent {
		keys = append(keys, currentName)
	}

	sort.Slice(keys, func(i, j int) bool {
		if keys[i] == currentName {
			return true
		}
		if keys[j] == currentName {
			return false
		}

		leftSpec := firstNodeSpecVariant(nodeSpecMap[keys[i]], withLocalDisk)
		rightSpec := firstNodeSpecVariant(nodeSpecMap[keys[j]], withLocalDisk)
		if leftSpec.GetCpu() != rightSpec.GetCpu() {
			return leftSpec.GetCpu() < rightSpec.GetCpu()
		}
		return leftSpec.GetMemoryGi() < rightSpec.GetMemoryGi()
	})

	candidates := make([]*apiv1.SharedPoolScaleConstraintNodeSpecCandidate, 0, len(keys))
	for _, key := range keys {
		specs := nodeSpecMap[key]
		candidateSpec, hasValid := resolveCandidateNodeSpec(specs, withLocalDisk)
		if candidateSpec == nil {
			candidateSpec = firstNodeSpecVariant(specs, withLocalDisk)
		}
		if candidateSpec == nil && key == currentName {
			candidateSpec = currentSpec
		}

		candidate := &apiv1.SharedPoolScaleConstraintNodeSpecCandidate{
			Name:      key,
			Cpu:       candidateSpec.GetCpu(),
			MemoryGi:  candidateSpec.GetMemoryGi(),
			IsCurrent: key == currentName,
		}

		reasons, available, candidateRequiredMinSpec := buildCandidateReasons(candidateSpec, hasValid, requiredMinSpec, storages)
		if candidate.IsCurrent {
			// Keep the current spec visible and selectable for UI continuity even if
			// it disappeared from the latest dataplane-derived candidate snapshot.
			// The paired scale flow is expected to omit node_spec when the user keeps
			// the current value; explicit resubmission of a disappeared spec is still
			// rejected by ScaleSharedPool today.
			available = true
			reasons = nil
			candidateRequiredMinSpec = nil
		}

		candidate.Available = available
		candidate.Reasons = reasons
		candidate.RequiredMinSpec = candidateRequiredMinSpec
		if available {
			candidate.StorageConstraints = buildStorageConstraints(storages, candidateSpec)
		}

		candidates = append(candidates, candidate)
	}

	return &apiv1.SharedPoolScaleConstraintNodeSpecConstraints{
		Current:    currentName,
		Candidates: candidates,
	}
}

func hasNodeSpecVariant(specs []*entity.NodeSpec, withLocalDisk bool) bool {
	for _, spec := range specs {
		if spec.GetWithLocalDisk() == withLocalDisk {
			return true
		}
	}
	return false
}

func firstNodeSpecVariant(specs []*entity.NodeSpec, withLocalDisk bool) *entity.NodeSpec {
	for _, spec := range specs {
		if spec.GetWithLocalDisk() == withLocalDisk {
			return spec
		}
	}
	return nil
}

func resolveCandidateNodeSpec(specs []*entity.NodeSpec, withLocalDisk bool) (*entity.NodeSpec, bool) {
	for _, spec := range specs {
		if spec.GetWithLocalDisk() != withLocalDisk {
			continue
		}
		if spec.GetAvailableCpu() == 0 || spec.GetAvailableMemoryMi() == 0 {
			continue
		}
		if spec.GetWithLocalDisk() && spec.GetLocalDiskSizeGi() == 0 {
			continue
		}
		return spec, true
	}
	return nil, false
}

func buildCandidateReasons(
	spec *entity.NodeSpec,
	hasValidSpec bool,
	requiredMinSpec *entity.NodeSpec,
	storages []*entity.Storage,
) ([]*apiv1.SharedPoolScaleConstraintReason, bool, *apiv1.SharedPoolScaleConstraintRequiredMinSpec) {
	reasons := make([]*apiv1.SharedPoolScaleConstraintReason, 0)
	available := hasValidSpec

	if !hasValidSpec {
		reasons = append(reasons, &apiv1.SharedPoolScaleConstraintReason{
			Code:    scaleConstraintReasonSpecNotPresent,
			Message: "target node spec is not currently available in shared pool",
		})
	}

	if requiredMinSpec != nil && isNodeSpecBelowRequired(spec, requiredMinSpec) {
		available = false
		reasons = append(reasons, &apiv1.SharedPoolScaleConstraintReason{
			Code:    scaleConstraintReasonBelowRURequirement,
			Message: "component node spec cannot support current minimum RU",
		})
	}

	if exceedsLocalDiskLimit(storages, spec) {
		available = false
		reasons = append(reasons, &apiv1.SharedPoolScaleConstraintReason{
			Code:    scaleConstraintReasonInsufficientLocalDisk,
			Message: "component local storage size exceeds target node spec local disk size",
		})
	}

	if requiredMinSpec != nil && slices.ContainsFunc(reasons, func(reason *apiv1.SharedPoolScaleConstraintReason) bool {
		return reason.GetCode() == scaleConstraintReasonBelowRURequirement
	}) {
		return reasons, available, &apiv1.SharedPoolScaleConstraintRequiredMinSpec{
			Name:     entity.GetComponentNodeSpec(requiredMinSpec),
			Cpu:      requiredMinSpec.GetCpu(),
			MemoryGi: requiredMinSpec.GetMemoryGi(),
		}
	}

	return reasons, available, nil
}

func getRequiredMinSpecForComponent(componentType apiv1.ComponentType, ruConfig *entity.RuConfig, minRU uint32) (*entity.NodeSpec, bool, error) {
	if minRU == 0 {
		return nil, false, nil
	}

	requiredMinSpec, err := ruConfig.GetNodeSpec(componentType, minRU)
	if err != nil {
		return nil, false, err
	}
	return requiredMinSpec, true, nil
}

func isNodeSpecBelowRequired(spec, requiredMinSpec *entity.NodeSpec) bool {
	if spec == nil || requiredMinSpec == nil {
		return false
	}
	return spec.GetCpu() < requiredMinSpec.GetCpu() || spec.GetMemoryGi() < requiredMinSpec.GetMemoryGi()
}

func exceedsLocalDiskLimit(storages []*entity.Storage, spec *entity.NodeSpec) bool {
	if spec == nil || !spec.GetWithLocalDisk() {
		return false
	}

	for _, storage := range storages {
		if storage.GetName() != apiv1.Component_Storage_DATA || storage.GetType() != entity.LocalSSD {
			continue
		}
		return storage.GetSizeGi() > spec.GetLocalDiskSizeGi()
	}

	return false
}

func buildStorageConstraints(storages []*entity.Storage, spec *entity.NodeSpec) []*apiv1.SharedPoolScaleConstraintStorageConstraint {
	out := make([]*apiv1.SharedPoolScaleConstraintStorageConstraint, 0, len(storages))
	for _, storage := range storages {
		item := &apiv1.SharedPoolScaleConstraintStorageConstraint{
			Name:          storage.GetName(),
			CurrentSizeGi: storage.GetSizeGi(),
		}

		switch {
		case storage.GetName() == apiv1.Component_Storage_DATA && storage.GetType() == entity.LocalSSD && spec.GetWithLocalDisk():
			size := spec.GetLocalDiskSizeGi()
			item.Resizable = false
			item.CurrentSizeGi = size
			item.MaxSizeGi = &size
		case storage.GetName() == apiv1.Component_Storage_DATA || storage.GetName() == apiv1.Component_Storage_RAFTLOG:
			item.Resizable = true
		default:
			item.Resizable = false
		}

		out = append(out, item)
	}
	return out
}