package server

import (
	"context"
	"fmt"
	"sort"
	"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/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"
)

type batchScaleSharedPoolService interface {
	ScaleSharedPoolBatch(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error)
}

type batchScaleSharedPoolWithLockHeldService interface {
	ScaleSharedPoolBatchWithLockHeld(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error)
}

// ScaleSharedPoolBatch submits a single update-component DAG for multiple shared-pool component changes.
func ScaleSharedPoolBatch(ctx context.Context, svc Service, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error) {
	batchSvc, ok := svc.(batchScaleSharedPoolService)
	if !ok {
		return nil, status.Errorf(codes.Unimplemented, "batch scale shared pool is not supported")
	}
	return batchSvc.ScaleSharedPoolBatch(ctx, sharedPoolId, scaleComponents, force)
}

// ScaleSharedPoolBatchWithLockHeld submits a batch update when the caller already
// holds the sharedPool advisory lock. This avoids re-entering the same lock.
func ScaleSharedPoolBatchWithLockHeld(ctx context.Context, svc Service, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error) {
	batchSvc, ok := svc.(batchScaleSharedPoolWithLockHeldService)
	if !ok {
		return nil, status.Errorf(codes.Unimplemented, "batch scale shared pool with lock held is not supported")
	}
	return batchSvc.ScaleSharedPoolBatchWithLockHeld(ctx, sharedPoolId, scaleComponents, force)
}

func (s *serverImpl) ScaleSharedPoolBatch(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error) {
	return s.scaleSharedPoolBatch(ctx, sharedPoolId, scaleComponents, force, false)
}

func (s *serverImpl) ScaleSharedPoolBatchWithLockHeld(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error) {
	return s.scaleSharedPoolBatch(ctx, sharedPoolId, scaleComponents, force, true)
}

func (s *serverImpl) scaleSharedPoolBatch(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force, lockHeld bool) (*entity.SharedPool, error) {
	if err := s.validateBatchScaleRequest(ctx, sharedPoolId, scaleComponents); err != nil {
		return nil, err
	}

	var sharedPool *entity.SharedPool
	var dagExecId string
	var dagSubmitted bool

	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		var err error
		if !lockHeld {
			unlockFunc, err := s.lock(ctx, sharedPoolId)
			if err != nil {
				tilog.Error(ctx, "failed to lock sharedPool", zap.String("sharedPoolId", sharedPoolId), zap.Error(err))
				return err
			}
			defer func() {
				if unlockErr := unlockFunc(); unlockErr != nil {
					tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", sharedPoolId), zap.Error(unlockErr))
				}
			}()
		}

		sharedPool, err = s.prepareBatchScaleOperation(ctx, sharedPoolId, scaleComponents, force)
		if err != nil {
			return err
		}

		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)
		}

		optionsMap, componentIds, componentTypes, err := s.mergeBatchScaleComponents(ctx, sharedPool, scaleComponents, nodeSpecMap, force)
		if err != nil {
			return err
		}
		componentIds, componentTypes = canonicalizeComponentRefs(componentIds, componentTypes)
		tilog.Info(ctx, "req-2: batch scale decision prepared",
			zap.String("sharedPoolId", sharedPoolId),
			zap.String("requestedComponents", describeComponents(scaleComponents)),
			zap.String("effectiveComponentIds", strings.Join(componentIds, ",")),
			zap.String("effectiveComponentTypes", strings.Join(componentTypes, ",")),
			zap.String("effectiveComponents", describeComponentSpecsByIDs(sharedPool, componentIds)),
			zap.String("options", describeOptionsMap(optionsMap)),
			zap.Bool("force", force),
			zap.Bool("lockHeld", lockHeld))
		if len(componentIds) == 0 {
			tilog.Info(ctx, "batch scale request produced no effective component changes", zap.String("sharedPoolId", sharedPoolId))
			AutoscaleRcuBatchScaleTotal.WithLabelValues(sharedPoolId, "no_changes").Inc()
			return nil
		}

		dagExecId, dagSubmitted, err = s.submitScaleBatchDAG(ctx, sharedPool, componentIds, componentTypes, optionsMap, force)
		return err
	}); 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 sharedPool, nil
}

func (s *serverImpl) validateBatchScaleRequest(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component) error {
	if sharedPoolId == "" {
		return status.Errorf(codes.InvalidArgument, "sharedPool id is required")
	}
	if len(scaleComponents) == 0 {
		return status.Errorf(codes.InvalidArgument, "at least one scale component is required")
	}
	for _, scaleComponent := range scaleComponents {
		if err := checkScaleComponentValid(ctx, scaleComponent, s.cloudProvider); err != nil {
			return err
		}
	}
	return nil
}

func (s *serverImpl) prepareBatchScaleOperation(ctx context.Context, sharedPoolId string, scaleComponents []*entity.Component, force bool) (*entity.SharedPool, error) {
	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		if status.Code(err) != codes.NotFound {
			return nil, status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
		}
		return nil, status.Errorf(codes.NotFound, "sharedPool %s not found", sharedPoolId)
	}
	if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
		return nil, 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)
	}

	// Forced operations are an explicit recovery escape hatch and may proceed
	// while the TiKV remote-worker topology is still converging.
	if !force {
		for _, scaleComponent := range scaleComponents {
			if isTiKVRemoteWorkerGuardedComponentType(scaleComponent.GetComponentType()) {
				if err := s.guardTiKVRemoteWorkersTopologyStable(ctx, sharedPool); err != nil {
					return nil, err
				}
				break
			}
		}
	}

	seen := make(map[string]struct{}, len(scaleComponents))
	for _, scaleComponent := range scaleComponents {
		if scaleComponent.GetId() == "" {
			tilog.Warn(ctx, "scale component id is empty, will use default id")
			scaleComponent.Id = sharedPool.GetDefaultComponenId()
		}

		if err := validateNodeCountForComponent(ctx, scaleComponent, sharedPool); err != nil {
			return nil, err
		}

		targetKey := batchScaleTargetKey(scaleComponent)
		if _, ok := seen[targetKey]; ok {
			return nil, status.Errorf(codes.InvalidArgument, "duplicate scale target %s", targetKey)
		}
		seen[targetKey] = struct{}{}

		if err := checkSharedPoolCanBeModified(ctx, sharedPool, scaleComponent, force); err != nil {
			return nil, err
		}

		if s.cloudProvider == "aws" {
			if err := s.checkVolumeCanBeModified(ctx, sharedPool, scaleComponent); err != nil {
				return nil, err
			}
		}
	}

	return sharedPool, nil
}

func (s *serverImpl) mergeBatchScaleComponents(ctx context.Context, sharedPool *entity.SharedPool, scaleComponents []*entity.Component, nodeSpecMap map[string][]*entity.NodeSpec, force bool) (map[string][]string, []string, []string, error) {
	sharedPool.Phase = apiv1.SharedPool_MODIFYING

	optionsMap := make(map[string][]string, len(scaleComponents))
	componentIds := make([]string, 0, len(scaleComponents))
	componentTypes := make([]string, 0, len(scaleComponents))
	seenTargets := make(map[string]struct{}, len(scaleComponents))

	for _, scaleComponent := range scaleComponents {
		componentOptions, err := s.mergeScaleComponent(ctx, sharedPool, scaleComponent, nodeSpecMap, force)
		if err != nil {
			return nil, nil, nil, err
		}

		for key, values := range componentOptions {
			optionsMap[key] = append(optionsMap[key], values...)
		}

		updateKey := entity.GetUpdateComponentKey(scaleComponent)
		if len(optionsMap[updateKey]) == 0 {
			continue
		}

		targetKey := batchScaleTargetKey(scaleComponent)
		if _, ok := seenTargets[targetKey]; ok {
			continue
		}
		seenTargets[targetKey] = struct{}{}
		componentIds = append(componentIds, scaleComponent.GetId())
		componentTypes = append(componentTypes, entity.GetComponentTypeName(scaleComponent))
	}

	return optionsMap, componentIds, componentTypes, nil
}

func (s *serverImpl) submitScaleBatchDAG(ctx context.Context, sharedPool *entity.SharedPool, componentIds, componentTypes []string, optionsMap map[string][]string, force bool) (string, bool, error) {
	sharedPoolStr, err := protojson.Marshal(sharedPool)
	if err != nil {
		return "", false, status.Errorf(codes.Internal, "failed to marshal sharedPool: %s", err)
	}

	args := actions.ActionArgs{
		workflow.SharedPoolIdKey:   actions.ActionArgString(sharedPool.GetId()),
		workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
		workflow.ComponentIdsKey:   actions.ActionArgString(strings.Join(componentIds, ",")),
		workflow.ComponentTypesKey: actions.ActionArgString(strings.Join(componentTypes, ",")),
		workflow.ForceKey:          actions.ActionArgString(fmt.Sprintf("%t", force)),
		utils.TraceIdKey:           utils.GetTraceId(ctx),
	}

	for k, v := range optionsMap {
		args[k] = actions.ActionArgString(strings.Join(v, ","))
	}
	tilog.Info(ctx, "req-2: submitting batch scale DAG",
		zap.String("sharedPoolId", sharedPool.GetId()),
		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("force", force))

	dagExecId := utils.GeneratedDagExecId(ctx,
		fmt.Sprintf("scale-sharedpool-batch-%s-%s-%s",
			sharedPool.GetId(),
			utils.FnvHash(strings.Join(componentTypes, ",")+"/"+strings.Join(componentIds, ",")),
			utils.FnvHash(string(sharedPoolStr)),
		))

	if err := s.dagExecutor.Submit(ctx, dag.PlanUpdateComponent, args, dagExecId); err != nil {
		return "", false, status.Errorf(codes.Internal, "failed to submit scale shared pool dag: %s", err)
	}

	AutoscaleRcuBatchScaleTotal.WithLabelValues(sharedPool.GetId(), "submitted").Inc()

	return dagExecId, true, nil
}

func batchScaleTargetKey(scaleComponent *entity.Component) string {
	return fmt.Sprintf("%s/%s", entity.GetComponentTypeName(scaleComponent), scaleComponent.GetId())
}

type componentRef struct {
	id   string
	kind string
}

func canonicalizeComponentRefs(componentIDs, componentTypes []string) ([]string, []string) {
	if len(componentIDs) != len(componentTypes) || len(componentIDs) <= 1 {
		return componentIDs, componentTypes
	}

	refs := make([]componentRef, 0, len(componentIDs))
	for i := range componentIDs {
		refs = append(refs, componentRef{
			id:   componentIDs[i],
			kind: componentTypes[i],
		})
	}
	sort.Slice(refs, func(i, j int) bool {
		if refs[i].kind == refs[j].kind {
			return refs[i].id < refs[j].id
		}
		return refs[i].kind < refs[j].kind
	})

	stableIDs := make([]string, 0, len(refs))
	stableTypes := make([]string, 0, len(refs))
	for _, ref := range refs {
		stableIDs = append(stableIDs, ref.id)
		stableTypes = append(stableTypes, ref.kind)
	}
	return stableIDs, stableTypes
}