package server

import (
	"context"
	"fmt"
	"strings"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/samber/lo"
	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) SetSharedPoolComponentConfigOrVersion(
	ctx context.Context,
	sharedPoolId string,
	targetComponents []*entity.TargetComponent,
	scope entity.SetVersionConfigScope,
	mode entity.SetVersionConfigMode,
	force bool,
) error {
	if sharedPoolId == "" {
		return status.Errorf(codes.InvalidArgument, "sharedPool id is required")
	}

	if len(targetComponents) == 0 {
		return status.Errorf(codes.InvalidArgument, "target components is empty")
	}

	for _, comp := range targetComponents {
		if comp.GetComponentType() == apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
			return status.Errorf(codes.InvalidArgument, "component type is required")
		}
		if comp.GetComponentType() == apiv1.ComponentType_DATA_MIGRATION {
			return status.Errorf(codes.InvalidArgument, "DATA_MIGRATION must be handled by SetDMComponentConfigOrVersion")
		}
	}

	// Validate scope early to fail fast
	if scope == apiv1.SetSharedPoolComponentConfigVersionRequest_SCOPE_UNSPECIFIED {
		return status.Errorf(codes.InvalidArgument, "invalid scope: %s", scope.String())
	}

	updateStrategy := v1alpha1.ConfigUpdateStrategyRestart
	if mode == apiv1.SetSharedPoolComponentConfigVersionRequest_HOT_RELOAD {
		updateStrategy = v1alpha1.ConfigUpdateStrategyHotReload
	}

	tiProxyLegacyPodsByGroup, err := s.prepareTiProxyGracefulRestartForTargets(ctx, sharedPoolId, targetComponents, mode)
	if err != nil {
		return err
	}

	var dagExecId string
	var dagSubmitted bool
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		// check if the sharedPool is locked
		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 err := unlockFunc(); err != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", sharedPoolId), zap.Error(err))
			}
		}()

		// get shared first
		sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, 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", sharedPoolId)
		}
		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)
		}

		var (
			ops                              []string
			ensureNoConstraintsUpgradePolicy bool
			optionsMap                       = map[string][]string{}
		)
		switch scope {
		case apiv1.SetSharedPoolComponentConfigVersionRequest_CONFIG:
			ops = []string{string(tioperator.OpsUpdateConfig), string(tioperator.OpsUpdateUpdateStrategy)}
		case apiv1.SetSharedPoolComponentConfigVersionRequest_VERSION:
			ops = []string{string(tioperator.OpsUpdateVersion), string(tioperator.OpsUpdateImage)}
			ensureNoConstraintsUpgradePolicy = true
		case apiv1.SetSharedPoolComponentConfigVersionRequest_ALL:
			ops = []string{string(tioperator.OpsUpdateVersion), string(tioperator.OpsUpdateImage), string(tioperator.OpsUpdateConfig), string(tioperator.OpsUpdateUpdateStrategy)}
			ensureNoConstraintsUpgradePolicy = true
		default:
			return status.Errorf(codes.InvalidArgument, "invalid or unsupported scope: %s", scope.String())
		}

		for _, comp := range targetComponents {
			if !isTiKVRemoteWorkerGuardedComponentType(comp.GetComponentType()) {
				continue
			}
			if err := s.guardTiKVRemoteWorkersTopologyStable(ctx, sharedPool); err != nil {
				return err
			}
			break
		}

		var componentIds, componentTypes []string
		for _, comp := range targetComponents {
			if comp.GetId() == "" {
				tilog.Warn(ctx, "target component id is empty, will use default id")
				comp.Id = sharedPool.GetDefaultComponenId()
			}

			targetComponent := &entity.Component{
				Id:            comp.GetId(),
				ComponentType: comp.GetComponentType(),
			}

			// check if the target component is valid
			err = checkSharedPoolCanBeModified(ctx, sharedPool, targetComponent, force)
			if err != nil {
				return err
			}

			var currComponent *entity.Component
			// get the current component
			for _, comp := range sharedPool.GetComponentGroups() {
				if comp.GetId() == targetComponent.GetId() && comp.GetComponentType() == targetComponent.GetComponentType() {
					currComponent = comp
					break
				}
			}

			if currComponent == nil {
				return status.Errorf(codes.NotFound, "component %s of type %s not found in sharedPool %s",
					targetComponent.GetId(), entity.GetComponentTypeName(targetComponent), sharedPool.GetId())
			}

			var tidbClusterId string
			if targetComponent.GetComponentType() == apiv1.ComponentType_TIDB {
				tidbGroup, err := s.tidbGroupRepo.Load(ctx, targetComponent.GetId(), false)
				if err != nil {
					if status.Code(err) != codes.NotFound {
						return status.Errorf(codes.Internal, "failed to load tidb group: %s", err)
					}
					// not found, check if it's in creating
					return status.Errorf(codes.NotFound, "tidb group %s not found", targetComponent.GetId())
				}
				tidbClusterId = tidbGroup.GetClusterId()
			}

			err = workflow.FullComponentConfigVersion(ctx, s.clusterConfigClient, s.meteringConfig, sharedPool, currComponent, tidbClusterId)
			if err != nil {
				tilog.Error(ctx, "failed to full target component config", zap.String("componentId", targetComponent.GetId()), zap.String("componentType", entity.GetComponentTypeName(targetComponent)), zap.Error(err))
				return status.Errorf(codes.Internal, "failed to full target component config: %s", err)
			}

			optionsMap[entity.GetUpdateComponentKey(currComponent)] = ops
			componentIds = append(componentIds, targetComponent.GetId())
			componentTypes = append(componentTypes, entity.GetComponentTypeName(targetComponent))
		}

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

		// start scale cluster dag
		args := actions.ActionArgs{
			workflow.SharedPoolIdKey:   actions.ActionArgString(sharedPool.GetId()),
			workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
			workflow.ComponentIdsKey:   actions.ActionArgString(strings.Join(lo.Union(componentIds), ",")),
			workflow.ComponentTypesKey: actions.ActionArgString(strings.Join(lo.Union(componentTypes), ",")),
			workflow.UpdateStrategyKey: actions.ActionArgString(string(updateStrategy)),
		}
		if ensureNoConstraintsUpgradePolicy {
			args[workflow.EnsureNoConstraintsUpgradePolicyKey] = actions.ActionArgBool(true)
		}

		// build update options
		for k, v := range optionsMap {
			args[k] = actions.ActionArgString(strings.Join(v, ","))
		}

		dagExecId = utils.GeneratedDagExecId(ctx,
			fmt.Sprintf("set-component-config-version-%s-%s",
				sharedPool.GetId(),
				utils.FnvHash(string(sharedPoolStr)),
			))
		err = s.dagExecutor.Submit(ctx, dag.PlanUpdateComponent, args, dagExecId)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to submit set-component-config-version 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 err
	}

	for tiproxyGroupId, podNames := range tiProxyLegacyPodsByGroup {
		if len(podNames) == 0 {
			continue
		}
		if err := s.markTiProxyBackendsNoNewConnections(ctx, sharedPoolId, tiproxyGroupId, podNames); err != nil {
			return status.Errorf(codes.Internal, "failed to mark tiproxy backends no-new-connections after rollout start: %s", err)
		}
	}

	return nil
}