package server

import (
	"context"
	"fmt"
	"math"

	"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"
	workflowaction "github.com/tidbcloud/cluster-service-ng/domain/workflow/action"
	"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/runtime/executor"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"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"
	"google.golang.org/protobuf/proto"
)

func (s *serverImpl) SetSharedPoolPDDeploymentMode(
	ctx context.Context,
	sharedPoolID string,
	targetMode entity.PDDeploymentMode,
) (*entity.SharedPool, error) {
	if sharedPoolID == "" {
		return nil, status.Errorf(codes.InvalidArgument, "sharedPool id is required")
	}
	if targetMode != apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MICROSERVICE &&
		targetMode != apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC {
		return nil, status.Errorf(codes.InvalidArgument, "pd deployment mode %s is not supported", targetMode.String())
	}

	var result *entity.SharedPool
	var dagExecID string
	var dagSubmitted bool
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		unlockFunc, err := s.lock(ctx, sharedPoolID)
		if err != nil {
			return err
		}
		defer func() {
			_ = unlockFunc()
		}()

		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)
			}
			return status.Errorf(codes.NotFound, "sharedPool %s not found", sharedPoolID)
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return err
		}

		if err := s.mergeSharedPoolStatusFromDagTask(ctx, sharedPool); err != nil {
			return status.Errorf(codes.Internal, "failed to merge sharedPool status from dag task: %s", err)
		}
		if sharedPool.GetCapacityMode() != apiv1.CapacityMode_RCU_DRIVEN {
			return status.Errorf(codes.FailedPrecondition, "pd deployment mode switch is only supported for RCU-driven shared pools")
		}

		runningDags, err := s.getRunningDagIns(ctx, pdDeploymentModeDAGPlans(), []string{sharedPoolID}, nil, nil, nil)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get running dag instance: %s", err)
		}
		if len(runningDags) > 0 {
			sameTarget, err := hasSamePDDeploymentModeTarget(runningDags, targetMode)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to inspect running pd deployment mode dag: %s", err)
			}
			if sameTarget {
				sharedPool.PdDeploymentMode = targetMode
				sharedPool.Phase = apiv1.SharedPool_MODIFYING
				s.mergeRunningPDDeploymentStatusBestEffort(ctx, sharedPool)
				result = sharedPool
				return nil
			}
			return status.Errorf(codes.FailedPrecondition, "another pd deployment mode switch is running for shared pool %s", sharedPoolID)
		}
		if sharedPool.GetPhase() != apiv1.SharedPool_AVAILABLE {
			return status.Errorf(codes.FailedPrecondition, "shared pool is not available, current phase: %s", sharedPool.GetPhase().String())
		}

		pdGroup, err := s.getSharedPoolPDGroup(ctx, sharedPool)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to read PDGroup actual mode status: %s", err)
		}
		// PDModeNormal is encoded as an empty string in the Operator API, so an
		// old monolithic PDGroup with empty status.mode is indistinguishable from
		// a new Operator reporting actual monolithic mode. Only the legacy
		// microservice shape is unambiguous: spec.mode=ms with status.mode still
		// empty means actual mode has not been populated yet, so switching must
		// fail fast instead of starting a DAG that cannot wait on actual status.
		if pdGroup.Spec.Template.Spec.Mode == tioperator.APIToPDMode(apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MICROSERVICE) &&
			pdGroup.Status.Mode == tioperator.APIToPDMode(apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC) {
			return status.Errorf(codes.FailedPrecondition, "PDGroup actual mode status is not ready; wait for TiDB Operator to populate PDGroup.status.mode before switching PD deployment mode")
		}
		if tioperator.PDModeToAPI(pdGroup.Status.Mode) == apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_UNSPECIFIED {
			return status.Errorf(codes.FailedPrecondition, "PDGroup actual mode status %q is not supported", pdGroup.Status.Mode)
		}

		pdStatus := buildPDDeploymentStatus(sharedPool, pdGroup, false)
		if isPDDeploymentModeStableInTarget(pdGroup, targetMode) {
			tempSharedPool, err := s.buildPDDeploymentModeTargetSharedPool(ctx, sharedPool, targetMode)
			if err != nil {
				return err
			}
			finalizeStablePDDeploymentModeSharedPool(sharedPool, tempSharedPool, targetMode)
			sharedPool.PdDeploymentStatus = buildPDDeploymentStatus(sharedPool, pdGroup, false)
			if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
				return fmt.Errorf("failed to save sharedPool meta: %w", err)
			}
			result = sharedPool
			return nil
		}

		tempSharedPool, err := s.buildPDDeploymentModeTargetSharedPool(ctx, sharedPool, targetMode)
		if err != nil {
			return err
		}

		sharedPool.PdDeploymentMode = targetMode
		if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
			return fmt.Errorf("failed to save sharedPool meta: %w", err)
		}

		runningSharedPool := buildPDDeploymentModeRunningSharedPool(sharedPool, tempSharedPool, targetMode)
		planName := pdDeploymentModePlanForTarget(targetMode)
		args, err := buildPDDeploymentModeArgs(ctx, runningSharedPool, targetMode)
		if err != nil {
			return err
		}
		dagExecID = utils.GeneratedDagExecId(ctx, fmt.Sprintf("set-pd-deployment-mode-%s-%s", sharedPoolID, pdDeploymentModeExecIDModeName(targetMode)))
		if err := s.dagExecutor.Submit(ctx, planName, args, dagExecID); err != nil {
			return status.Errorf(codes.Internal, "failed to submit pd deployment mode dag: %s", err)
		}
		dagSubmitted = true

		runningSharedPool.PdDeploymentStatus = buildRunningPDDeploymentStatus(pdStatus)
		result = runningSharedPool
		return nil
	}); err != nil {
		if dagSubmitted && dagExecID != "" {
			if abortErr := s.dagExecutor.Abort(ctx, dagExecID); abortErr != nil {
				tilog.Error(ctx, "failed to abort pd deployment mode dag after transaction failure",
					zap.String("dagExecID", dagExecID),
					zap.Error(abortErr))
			}
		}
		return nil, err
	}

	return result, nil
}

func (s *serverImpl) buildPDDeploymentModeTargetSharedPool(
	ctx context.Context,
	sharedPool *entity.SharedPool,
	targetMode entity.PDDeploymentMode,
) (*entity.SharedPool, error) {
	tempSharedPool := &entity.SharedPool{
		TemplateVersion:           sharedPool.TemplateVersion,
		TenantId:                  sharedPool.TenantId,
		DeployType:                sharedPool.DeployType,
		ConfigCenterMetaProfileId: sharedPool.ConfigCenterMetaProfileId,
		CreatedAt:                 sharedPool.CreatedAt,
		DeletedAt:                 sharedPool.DeletedAt,
	}
	proto.Merge(&tempSharedPool.SharedPool, &sharedPool.SharedPool)
	tempSharedPool.PdDeploymentMode = targetMode
	tempSharedPool.Phase = apiv1.SharedPool_MODIFYING

	switch targetMode {
	case apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MICROSERVICE:
		ruConfig, err := s.ruConfigRepo.GetRuConfig(ctx, sharedPool.GetId(), apiv1.RuConfig_SHARED_POOL)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to get shared pool ru config: %s", err)
		}
		nodeSpecMap, err := s.genNodeSpecMap(ctx, sharedPool)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to get node spec map: %s", err)
		}
		ruTotal, err := s.getSharedPoolRequestUnitTotal(ctx, sharedPool.GetId())
		if err != nil {
			return nil, err
		}
		defaultComponentMap := sharedPoolDefaultComponentMap(tempSharedPool)
		if err := s.generatedSharedPoolMicroservicePDComponent(ctx, tempSharedPool, defaultComponentMap, ruTotal, nodeSpecMap, ruConfig, false); err != nil {
			return nil, err
		}
		markComponentsPhase(tempSharedPool, microservicePDComponentTypes(), apiv1.Component_MODIFYING)
	case apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC:
		markComponentsPhase(tempSharedPool, []entity.ComponentType{apiv1.ComponentType_PD}, apiv1.Component_MODIFYING)
		markComponentsPhase(tempSharedPool, microservicePDComponentTypes(), apiv1.Component_DELETING)
	}
	return tempSharedPool, nil
}

func buildPDDeploymentModeRunningSharedPool(
	sharedPool *entity.SharedPool,
	targetSharedPool *entity.SharedPool,
	targetMode entity.PDDeploymentMode,
) *entity.SharedPool {
	runningSharedPool := &entity.SharedPool{
		TemplateVersion:           sharedPool.TemplateVersion,
		TenantId:                  sharedPool.TenantId,
		DeployType:                sharedPool.DeployType,
		ConfigCenterMetaProfileId: sharedPool.ConfigCenterMetaProfileId,
		CreatedAt:                 sharedPool.CreatedAt,
		DeletedAt:                 sharedPool.DeletedAt,
	}
	proto.Merge(&runningSharedPool.SharedPool, &sharedPool.SharedPool)
	runningSharedPool.ComponentGroups = workflowaction.MergePDDeploymentModeComponents(
		sharedPool.GetComponentGroups(),
		targetSharedPool.GetComponentGroups(),
	)
	runningSharedPool.PdDeploymentMode = targetMode
	runningSharedPool.Phase = targetSharedPool.GetPhase()
	return runningSharedPool
}

func finalizeStablePDDeploymentModeSharedPool(
	sharedPool *entity.SharedPool,
	targetSharedPool *entity.SharedPool,
	targetMode entity.PDDeploymentMode,
) {
	sharedPool.ComponentGroups = workflowaction.MergePDDeploymentModeComponents(
		sharedPool.GetComponentGroups(),
		targetSharedPool.GetComponentGroups(),
	)
	workflowaction.FinalizePDDeploymentModeSharedPool(sharedPool, targetMode)
}

func (s *serverImpl) getSharedPoolRequestUnitTotal(ctx context.Context, sharedPoolID string) (*entity.RequestUnit, error) {
	tidbClusters, err := s.tidbClusterRepo.List(ctx, &repo.TiDBClusterListFilter{
		SharedPoolIds: []string{sharedPoolID},
		Deleted:       false,
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list tidb cluster: %s", err)
	}

	var ruMaxTotal, ruMinTotal uint64
	for _, tidbCluster := range tidbClusters {
		ruMaxTotal += uint64(tidbCluster.GetRequestUnit().GetMax())
		ruMinTotal += uint64(tidbCluster.GetRequestUnit().GetMin())
	}
	if ruMaxTotal > math.MaxUint32 || ruMinTotal > math.MaxUint32 {
		return nil, status.Errorf(codes.InvalidArgument, "total RU overflows uint32")
	}
	return &entity.RequestUnit{
		Max: uint32(ruMaxTotal),
		Min: uint32(ruMinTotal),
	}, nil
}

func buildPDDeploymentModeArgs(
	ctx context.Context,
	sharedPool *entity.SharedPool,
	targetMode entity.PDDeploymentMode,
) (actions.ActionArgs, error) {
	sharedPoolStr, err := protojson.Marshal(sharedPool)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to marshal sharedPool: %s", err)
	}

	return actions.ActionArgs{
		workflow.SharedPoolIdKey:     actions.ActionArgString(sharedPool.GetId()),
		workflow.SharedPoolKey:       actions.ActionArgString(string(sharedPoolStr)),
		workflow.ComponentIdsKey:     actions.ActionArgString(defaultSharedPoolComponentID(sharedPool)),
		workflow.PDDeploymentModeKey: actions.ActionArgString(targetMode.String()),
		utils.TraceIdKey:             utils.GetTraceId(ctx),
	}, nil
}

func hasSamePDDeploymentModeTarget(runningDags executor.ExecTaskIdStatus, targetMode entity.PDDeploymentMode) (bool, error) {
	for _, runningDag := range runningDags {
		modeStr, err := runningDag.ExecTask.Args.GetString(workflow.PDDeploymentModeKey)
		if err != nil {
			return false, err
		}
		if modeStr == targetMode.String() {
			return true, nil
		}
	}
	return false, nil
}

func pdDeploymentModeDAGPlans() []string {
	return []string{
		dag.PlanSetSharedPoolPDDeploymentModeToMonolithic,
		dag.PlanSetSharedPoolPDDeploymentModeToMicroservice,
	}
}

func pdDeploymentModePlanForTarget(targetMode entity.PDDeploymentMode) string {
	if targetMode == apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC {
		return dag.PlanSetSharedPoolPDDeploymentModeToMonolithic
	}
	return dag.PlanSetSharedPoolPDDeploymentModeToMicroservice
}

func pdDeploymentModeExecIDModeName(targetMode entity.PDDeploymentMode) string {
	if targetMode == apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC {
		return "monolithic"
	}
	return "microservice"
}

func sharedPoolDefaultComponentMap(sharedPool *entity.SharedPool) map[entity.ComponentType]*entity.Component {
	defaultComponentID := defaultSharedPoolComponentID(sharedPool)

	defaultComponentMap := make(map[entity.ComponentType]*entity.Component)
	for _, comp := range sharedPool.GetComponentGroups() {
		if comp.GetId() == defaultComponentID {
			defaultComponentMap[comp.GetComponentType()] = comp
		}
	}
	return defaultComponentMap
}

func markComponentsPhase(sharedPool *entity.SharedPool, componentTypes []entity.ComponentType, phase entity.ComponentPhase) {
	defaultComponentID := defaultSharedPoolComponentID(sharedPool)
	for _, comp := range sharedPool.GetComponentGroups() {
		if comp.GetId() == defaultComponentID && lo.Contains(componentTypes, comp.GetComponentType()) {
			comp.Phase = phase
		}
	}
}

func defaultSharedPoolComponentID(sharedPool *entity.SharedPool) string {
	if defaultComponentID := sharedPool.GetDefaultComponenId(); defaultComponentID != "" {
		return defaultComponentID
	}
	return utils.FnvHash(sharedPool.GetId())
}

func microservicePDComponentTypes() []entity.ComponentType {
	return []entity.ComponentType{
		apiv1.ComponentType_PD_TSO,
		apiv1.ComponentType_PD_SCHEDULING,
		apiv1.ComponentType_RESOURCE_MANAGER,
	}
}