package server

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"unicode"

	"github.com/google/uuid"
	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/dataplane"
	"github.com/tidbcloud/cluster-service-ng/pkg/nodegroup"
	"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/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"
)

// Component label keys used by the cassandra-proxy handler. The action that
// renders the K8s manifests reads these same keys, so renaming requires a
// coordinated update on the action side (apply_cassandra_proxy.go).
const (
	cassandraProxyHostIdLabelKey   = "cassandra-proxy/host-id"
	cassandraProxyNodeSpecLabelKey = "cassandra-proxy/node-spec"
)

// errCassandraProxyComponentExists is a sentinel that signals a concurrent
// Set raced ahead of us inside the locked handleCreate path; the outer
// dispatcher catches this and re-routes to handleUpdate.
var errCassandraProxyComponentExists = errors.New("cassandra-proxy component already exists")

// SetCassandraProxy is the declarative desired-state RPC for the
// cassandra-proxy component. Branches on whether a component already exists
// in the SP; new -> handleCreate (append + DAG), existing -> handleUpdate
// (phase-driven no-op / scale stub).
func (s *serverImpl) SetCassandraProxy(ctx context.Context, req *apiv1.SetCassandraProxyRequest) (*apiv1.SetCassandraProxyResponse, error) {
	if req.GetClusterId() == "" {
		return nil, status.Errorf(codes.InvalidArgument, "cluster_id is required")
	}
	nodeSpec, err := cassandraProxyNodeSpecFromQuota(req.GetNodeQuota())
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "%s", err)
	}

	tidbCluster, err := s.tidbClusterRepo.Load(ctx, req.GetClusterId(), false)
	if err != nil {
		if status.Code(err) == codes.NotFound {
			return nil, status.Errorf(codes.NotFound, "tidb cluster %s not found", req.GetClusterId())
		}
		return nil, status.Errorf(codes.Internal, "failed to load tidb cluster: %s", err)
	}

	sharedPoolId := tidbCluster.GetSharedPoolId()
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.FailedPrecondition, "tidb cluster %s has no shared_pool_id", req.GetClusterId())
	}
	if tidbCluster.GetGatewayIdentity() == "" {
		return nil, status.Errorf(codes.FailedPrecondition,
			"tidb cluster %s has empty gateway_identity; required for Cassandra proxy PrivateLink LB",
			req.GetClusterId())
	}

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

	// PrivateLink is unconditional for cassandra-proxy, so the dataplane must
	// support multi-port LB specs and annotation passthrough. Probe before any
	// state-changing work so unsupported dataplanes fail the RPC early instead
	// of leaving a half-applied component behind.
	if err := dataplane.RequireDataplaneCapability(
		ctx,
		s.dataplaneClient,
		sharedPool.GetDataplaneName(),
		dataplane.CapabilityLoadBalancerMultiPort,
	); err != nil {
		return nil, err
	}

	state, err := s.cassandraProxyStateRepo.LoadByClusterID(ctx, req.GetClusterId(), false)
	if err != nil && status.Code(err) != codes.NotFound {
		return nil, status.Errorf(codes.Internal, "failed to load cassandra-proxy state: %s", err)
	}
	if state != nil {
		return s.handleCassandraProxyUpdate(ctx, req, state, nodeSpec)
	}
	resp, err := s.handleCassandraProxyCreate(ctx, sharedPoolId, req, nodeSpec)
	if errors.Is(err, errCassandraProxyComponentExists) || status.Code(err) == codes.AlreadyExists {
		// concurrent create raced us; reload persisted state and dispatch as update.
		state, loadErr := s.cassandraProxyStateRepo.LoadByClusterID(ctx, req.GetClusterId(), false)
		if loadErr != nil {
			return nil, status.Errorf(codes.Internal, "failed to reload cassandra-proxy state after race: %s", loadErr)
		}
		if state == nil {
			return nil, status.Errorf(codes.Internal, "cassandra-proxy state vanished after create race for cluster %s", req.GetClusterId())
		}
		return s.handleCassandraProxyUpdate(ctx, req, state, nodeSpec)
	}
	return resp, err
}

func cassandraProxyNodeSpecFromQuota(nodeQuota *apiv1.NodeQuota) (string, error) {
	if nodeQuota == nil {
		return "", fmt.Errorf("node_quota is required")
	}
	nodeSpec := nodeQuota.GetNodeSpec()
	if nodeSpec == nil {
		return "", fmt.Errorf("node_quota.node_spec is required")
	}
	if nodeSpec.GetCpu() < 2 {
		return "", fmt.Errorf("node_quota.node_spec.cpu must be >= 2")
	}
	if nodeSpec.GetMemoryGi() < 4 {
		return "", fmt.Errorf("node_quota.node_spec.memory_gi must be >= 4")
	}
	value := entity.GetComponentNodeSpec(nodeSpec)
	if err := nodegroup.ValidateSpec(value); err != nil {
		return "", fmt.Errorf("invalid node_quota.node_spec: %w", err)
	}
	return value, nil
}

func (s *serverImpl) configuredCassandraProxyImage() (string, error) {
	if s.imageConfig == nil {
		return "", status.Errorf(codes.FailedPrecondition, "cassandra-proxy image is not configured (imageConfig.cassandraProxyImage)")
	}
	image := strings.TrimSpace(s.imageConfig.CassandraProxyImage)
	if image == "" {
		return "", status.Errorf(codes.FailedPrecondition, "cassandra-proxy image is not configured (imageConfig.cassandraProxyImage)")
	}
	if err := validateCassandraProxyImage(image); err != nil {
		return "", status.Errorf(codes.FailedPrecondition, "invalid imageConfig.cassandraProxyImage: %s", err)
	}
	return image, nil
}

func validateCassandraProxyImage(image string) error {
	if strings.TrimSpace(image) == "" {
		return fmt.Errorf("image is empty")
	}
	if len(image) > 512 {
		return fmt.Errorf("image exceeds 512 characters")
	}
	if strings.ContainsFunc(image, func(r rune) bool {
		return unicode.IsSpace(r) || unicode.IsControl(r)
	}) {
		return fmt.Errorf("image must not contain whitespace or control characters")
	}
	return nil
}

func (s *serverImpl) resolveCassandraProxyImageForCreate(req *apiv1.SetCassandraProxyRequest) (string, error) {
	if req.Image == nil {
		return s.configuredCassandraProxyImage()
	}
	image := req.GetImage()
	if err := validateCassandraProxyImage(image); err != nil {
		return "", status.Errorf(codes.InvalidArgument, "invalid image: %s", err)
	}
	return image, nil
}

func (s *serverImpl) resolveCassandraProxyImageFromState(state *entity.CassandraProxyState) (string, error) {
	image := strings.TrimSpace(state.Image)
	if image == "" {
		return s.configuredCassandraProxyImage()
	}
	if err := validateCassandraProxyImage(image); err != nil {
		return "", status.Errorf(codes.FailedPrecondition, "invalid cassandra-proxy state image: %s", err)
	}
	return image, nil
}

func (s *serverImpl) resolveCassandraProxyImageForUpdate(req *apiv1.SetCassandraProxyRequest, state *entity.CassandraProxyState) (string, bool, error) {
	if req.Image != nil {
		image := req.GetImage()
		if err := validateCassandraProxyImage(image); err != nil {
			return "", false, status.Errorf(codes.InvalidArgument, "invalid image: %s", err)
		}
		return image, image != state.Image, nil
	}
	image, err := s.resolveCassandraProxyImageFromState(state)
	if err != nil {
		return "", false, err
	}
	return image, strings.TrimSpace(state.Image) != "" && image != state.Image, nil
}

func (s *serverImpl) handleCassandraProxyCreate(ctx context.Context, sharedPoolId string, req *apiv1.SetCassandraProxyRequest, nodeSpec string) (*apiv1.SetCassandraProxyResponse, error) {
	image, err := s.resolveCassandraProxyImageForCreate(req)
	if err != nil {
		return nil, err
	}

	var (
		hostId       string
		dagExecId    string
		dagSubmitted bool
		respPhase    apiv1.Component_Phase
		respEndpoint string
	)

	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		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))
			}
		}()

		sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
		if err != nil {
			if status.Code(err) == codes.NotFound {
				return status.Errorf(codes.NotFound, "shared pool %s not found", sharedPoolId)
			}
			return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
		}
		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 err := checkSharedPoolCanAcceptNewComponent(sharedPool); err != nil {
			return err
		}

		// TOCTOU recheck: if a concurrent create slipped in, return sentinel
		// so the outer dispatcher routes to handleUpdate.
		state, err := s.cassandraProxyStateRepo.LoadByClusterID(ctx, req.GetClusterId(), false)
		if err != nil && status.Code(err) != codes.NotFound {
			return status.Errorf(codes.Internal, "failed to load cassandra-proxy state: %s", err)
		}
		if state != nil {
			return errCassandraProxyComponentExists
		}

		hostId = uuid.NewString()
		state = &entity.CassandraProxyState{
			ClusterID:          req.GetClusterId(),
			SharedPoolID:       sharedPool.GetId(),
			HostID:             hostId,
			NodeSpec:           nodeSpec,
			Image:              image,
			Phase:              entity.ComponentPhaseToCassandraProxyPhase(apiv1.Component_CREATING),
			PrivatelinkDNSName: "",
			NodeQuotaEntity:    req.GetNodeQuota(),
		}
		component := newCassandraProxyComponent(state, apiv1.Component_CREATING)
		sharedPool.ComponentGroups = removeComponentByIDAndType(
			sharedPool.ComponentGroups,
			req.GetClusterId(),
			apiv1.ComponentType_CASSANDRA_PROXY,
		)
		sharedPool.ComponentGroups = append(sharedPool.ComponentGroups, component)
		sharedPool.Phase = apiv1.SharedPool_MODIFYING
		if err := s.cassandraProxyStateRepo.Save(ctx, state); err != nil {
			if status.Code(err) == codes.AlreadyExists {
				return errCassandraProxyComponentExists
			}
			return status.Errorf(codes.Internal, "failed to save cassandra-proxy state: %s", err)
		}

		// Save the shared pool with the Cassandra proxy component so the frozen
		// workflow.SharedPoolKey snapshot below has the desired apply input. Durable
		// phase reads come from cassandraProxyStateRepo, not SharedPool.ComponentGroups.
		if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
			return status.Errorf(codes.Internal, "failed to save sharedPool: %s", err)
		}

		sharedPoolStr, err := protojson.Marshal(sharedPool)
		if err != nil {
			return 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(req.GetClusterId()),
			workflow.ComponentTypesKey: actions.ActionArgString(entity.GetComponentTypeName(component)),
			workflow.TiDBClusterIdKey:  actions.ActionArgString(req.GetClusterId()),
			utils.TraceIdKey:           utils.GetTraceId(ctx),
		}

		dagExecId = newCreateCassandraProxyDagExecID(ctx, sharedPool.GetId(), req.GetClusterId(), string(sharedPoolStr), "")

		if err := s.dagExecutor.Submit(ctx, dag.PlanCreateCassandraProxy, args, dagExecId); err != nil {
			return status.Errorf(codes.Internal, "failed to submit create-cassandra-proxy dag: %s", err)
		}
		dagSubmitted = true

		respPhase = apiv1.Component_CREATING
		respEndpoint = buildCassandraProxyInternalEndpoint(sharedPool.GetId(), req.GetClusterId())

		tilog.Info(ctx, "submitted create cassandra-proxy dag",
			zap.String("sharedPoolId", sharedPoolId),
			zap.String("clusterId", req.GetClusterId()),
			zap.String("dagExecId", dagExecId))
		return nil
	}); err != nil {
		if errors.Is(err, errCassandraProxyComponentExists) {
			return nil, err
		}
		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 &apiv1.SetCassandraProxyResponse{
		HostId:             hostId,
		InternalEndpoint:   respEndpoint,
		PrivatelinkDnsName: "",
		Phase:              respPhase,
		Image:              image,
	}, nil
}

func (s *serverImpl) handleCassandraProxyUpdate(ctx context.Context, req *apiv1.SetCassandraProxyRequest, state *entity.CassandraProxyState, nodeSpec string) (*apiv1.SetCassandraProxyResponse, error) {
	phase := entity.CassandraProxyPhaseToComponentPhase(state.Phase)
	image, imageChanged, err := s.resolveCassandraProxyImageForUpdate(req, state)
	if err != nil {
		return nil, err
	}
	sizingChanged := cassandraProxySizingChanged(req, state, nodeSpec)

	switch phase {
	case apiv1.Component_CREATING, apiv1.Component_MODIFYING:
		// Reject spec changes while in-flight: silently accepting would leave the
		// DAG snapshot stale relative to the desired state the caller just sent.
		if sizingChanged || imageChanged {
			return nil, status.Errorf(codes.FailedPrecondition,
				"cassandra proxy is %s, retry after AVAILABLE", phase)
		}
		resp := buildCassandraProxySetResponse(state, phase, image)
		dagIns, err := s.getRunningDagIns(ctx,
			[]string{dag.PlanCreateCassandraProxy},
			[]string{state.SharedPoolID},
			nil,
			nil,
			[]string{state.ClusterID})
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to get running cassandra-proxy dag instance: %s", err)
		}
		if len(dagIns) == 0 {
			freshState, updated, err := s.cassandraProxyStateRepo.UpdatePhaseIfCurrent(ctx,
				state.ClusterID,
				state.HostID,
				[]entity.ComponentPhase{apiv1.Component_CREATING, apiv1.Component_MODIFYING},
				apiv1.Component_UNAVAILABLE,
			)
			if err != nil {
				return nil, status.Errorf(codes.Internal, "failed to mark cassandra-proxy unavailable: %s", err)
			}
			if !updated {
				freshPhase := entity.CassandraProxyPhaseToComponentPhase(freshState.Phase)
				freshImage, err := s.resolveCassandraProxyImageFromState(freshState)
				if err != nil {
					return nil, err
				}
				return buildCassandraProxySetResponse(freshState, freshPhase, freshImage), nil
			}
			return nil, status.Errorf(codes.FailedPrecondition, "previous cassandra-proxy create failed; SetCassandraProxy cannot retry from UNAVAILABLE")
		}
		// in-flight: idempotent no-op
		return resp, nil
	case apiv1.Component_DELETING:
		// out-of-band tear-down somehow ran; surface for ops attention.
		return nil, status.Errorf(codes.FailedPrecondition, "cassandra-proxy component is being deleted, retry later")
	case apiv1.Component_UNAVAILABLE:
		return s.retryUnavailableCassandraProxy(ctx, req, state, nodeSpec, imageChanged)
	case apiv1.Component_AVAILABLE:
		if sizingChanged || imageChanged {
			return s.handleCassandraProxyScale(ctx, req, state, nodeSpec)
		}
		return s.buildCassandraProxySetResponse(ctx, state, phase, image), nil
	}
	return buildCassandraProxySetResponse(state, phase, image), nil
}

func (s *serverImpl) retryUnavailableCassandraProxy(ctx context.Context, req *apiv1.SetCassandraProxyRequest, state *entity.CassandraProxyState, nodeSpec string, imageChanged bool) (*apiv1.SetCassandraProxyResponse, error) {
	if (cassandraProxySizingChanged(req, state, nodeSpec) || imageChanged) && !req.GetForce() {
		return nil, status.Errorf(codes.FailedPrecondition, "cannot change cassandra-proxy desired state while component is UNAVAILABLE; retry with the current desired state first")
	}

	var (
		resp         *apiv1.SetCassandraProxyResponse
		dagExecId    string
		dagSubmitted bool
	)
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		unlockFunc, err := s.lock(ctx, state.SharedPoolID)
		if err != nil {
			tilog.Error(ctx, "failed to lock sharedPool", zap.String("sharedPoolId", state.SharedPoolID), zap.Error(err))
			return err
		}
		defer func() {
			if err := unlockFunc(); err != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", state.SharedPoolID), zap.Error(err))
			}
		}()

		freshState, err := s.cassandraProxyStateRepo.LoadByClusterID(ctx, state.ClusterID, false)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to reload cassandra-proxy state before retry: %s", err)
		}
		if freshState == nil {
			return status.Errorf(codes.Internal, "cassandra-proxy state vanished before retry for cluster %s", state.ClusterID)
		}

		phase := entity.CassandraProxyPhaseToComponentPhase(freshState.Phase)
		if phase == apiv1.Component_DELETING {
			return status.Errorf(codes.FailedPrecondition, "cassandra-proxy component is being deleted, retry later")
		}
		image, freshImageChanged, err := s.resolveCassandraProxyImageForUpdate(req, freshState)
		if err != nil {
			return err
		}
		freshSpecChanged := cassandraProxySizingChanged(req, freshState, nodeSpec) || freshImageChanged
		if phase != apiv1.Component_UNAVAILABLE {
			if freshSpecChanged {
				if phase == apiv1.Component_CREATING || phase == apiv1.Component_MODIFYING {
					return status.Errorf(codes.FailedPrecondition, "cassandra proxy is %s, retry after AVAILABLE", phase)
				}
				return status.Errorf(codes.FailedPrecondition, "cassandra-proxy state changed before applying desired state, retry request")
			}
			resp = buildCassandraProxySetResponse(freshState, phase, image)
			return nil
		}
		if freshSpecChanged && !req.GetForce() {
			return status.Errorf(codes.FailedPrecondition, "cannot change cassandra-proxy desired state while component is UNAVAILABLE; retry with the current desired state first")
		}
		if image == "" {
			image, err = s.configuredCassandraProxyImage()
			if err != nil {
				return err
			}
		}

		dagIns, err := s.getRunningDagIns(ctx,
			[]string{dag.PlanCreateCassandraProxy},
			[]string{freshState.SharedPoolID},
			nil,
			nil,
			[]string{freshState.ClusterID})
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get running cassandra-proxy dag instance: %s", err)
		}
		if freshSpecChanged && len(dagIns) > 0 {
			return status.Errorf(codes.FailedPrecondition, "cannot force change cassandra-proxy desired state while a workflow is running")
		}
		if hasRunningCassandraProxyCreateDagForState(ctx, dagIns, freshState) {
			// If executor cleanup races this query, the next Set call will mark the
			// CREATING state back to UNAVAILABLE and retry from there.
			updatedState, updated, err := s.updateUnavailableCassandraProxyToCreating(ctx, freshState)
			if err != nil {
				return err
			}
			if !updated {
				resp = buildCassandraProxySetResponse(updatedState, entity.CassandraProxyPhaseToComponentPhase(updatedState.Phase), image)
				return nil
			}
			resp = buildCassandraProxySetResponse(updatedState, apiv1.Component_CREATING, image)
			return nil
		}
		sharedPool, err := s.sharedPoolRepo.Load(ctx, freshState.SharedPoolID, false)
		if err != nil {
			if status.Code(err) == codes.NotFound {
				return status.Errorf(codes.NotFound, "shared pool %s not found", freshState.SharedPoolID)
			}
			return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
		}
		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 err := checkSharedPoolCanAcceptNewComponent(sharedPool); err != nil && (!req.GetForce() || sharedPool.GetPhase() != apiv1.SharedPool_UNAVAILABLE) {
			return err
		}

		stateForDag := *freshState
		stateForDag.Image = image
		if freshSpecChanged {
			stateForDag.NodeSpec = nodeSpec
			stateForDag.NodeQuotaEntity = req.GetNodeQuota()
		}
		component := newCassandraProxyComponent(&stateForDag, apiv1.Component_CREATING)
		sharedPool.ComponentGroups = removeComponentByIDAndType(
			sharedPool.ComponentGroups,
			freshState.ClusterID,
			apiv1.ComponentType_CASSANDRA_PROXY,
		)
		sharedPool.ComponentGroups = append(sharedPool.ComponentGroups, component)
		sharedPool.Phase = apiv1.SharedPool_MODIFYING
		updatedState, updated, err := s.updateUnavailableCassandraProxyToCreating(ctx, freshState)
		if err != nil {
			return err
		}
		if !updated {
			if freshSpecChanged {
				return status.Errorf(codes.FailedPrecondition, "cassandra-proxy state changed before applying desired state, retry request")
			}
			resp = buildCassandraProxySetResponse(updatedState, entity.CassandraProxyPhaseToComponentPhase(updatedState.Phase), image)
			return nil
		}
		if updatedState.Image == "" || freshSpecChanged {
			updatedState.Image = stateForDag.Image
			updatedState.NodeSpec = stateForDag.NodeSpec
			updatedState.NodeQuotaEntity = stateForDag.NodeQuotaEntity
			if err := s.cassandraProxyStateRepo.Save(ctx, updatedState); err != nil {
				return status.Errorf(codes.Internal, "failed to save cassandra-proxy desired state before retry: %s", err)
			}
		}
		if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
			return status.Errorf(codes.Internal, "failed to save sharedPool: %s", err)
		}

		sharedPoolStr, err := protojson.Marshal(sharedPool)
		if err != nil {
			return 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(freshState.ClusterID),
			workflow.ComponentTypesKey: actions.ActionArgString(entity.GetComponentTypeName(component)),
			workflow.TiDBClusterIdKey:  actions.ActionArgString(freshState.ClusterID),
			utils.TraceIdKey:           utils.GetTraceId(ctx),
		}
		dagExecId = newCreateCassandraProxyDagExecID(ctx, sharedPool.GetId(), freshState.ClusterID, string(sharedPoolStr), uuid.NewString())
		if err := s.dagExecutor.Submit(ctx, dag.PlanCreateCassandraProxy, args, dagExecId); err != nil {
			return status.Errorf(codes.Internal, "failed to submit create-cassandra-proxy dag: %s", err)
		}
		dagSubmitted = true
		resp = buildCassandraProxySetResponse(freshState, apiv1.Component_CREATING, image)
		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 nil, err
	}
	return resp, nil
}

func (s *serverImpl) updateUnavailableCassandraProxyToCreating(ctx context.Context, state *entity.CassandraProxyState) (*entity.CassandraProxyState, bool, error) {
	updatedState, updated, err := s.cassandraProxyStateRepo.UpdatePhaseIfCurrent(ctx,
		state.ClusterID,
		state.HostID,
		[]entity.ComponentPhase{apiv1.Component_UNAVAILABLE},
		apiv1.Component_CREATING,
	)
	if err != nil {
		return nil, false, status.Errorf(codes.Internal, "failed to mark cassandra-proxy creating: %s", err)
	}
	if updatedState == nil {
		return nil, false, status.Errorf(codes.Internal, "cassandra-proxy state vanished before marking creating for cluster %s", state.ClusterID)
	}
	if !updated && entity.CassandraProxyPhaseToComponentPhase(updatedState.Phase) == apiv1.Component_DELETING {
		return updatedState, false, status.Errorf(codes.FailedPrecondition, "cassandra-proxy component is being deleted, retry later")
	}
	return updatedState, updated, nil
}

func hasRunningCassandraProxyCreateDagForState(ctx context.Context, dagIns executor.ExecTaskIdStatus, state *entity.CassandraProxyState) bool {
	for _, dagIn := range dagIns {
		if dagIn == nil || dagIn.ExecTask == nil {
			continue
		}
		sharedPoolStr, err := dagIn.ExecTask.Args.GetString(workflow.SharedPoolKey)
		if err != nil {
			tilog.Warn(ctx, "hasRunningCassandraProxyCreateDagForState: failed to get shared pool key from dag args", zap.Error(err))
			continue
		}
		var sharedPool entity.SharedPool
		if err := protojson.Unmarshal([]byte(sharedPoolStr), &sharedPool); err != nil {
			tilog.Warn(ctx, "hasRunningCassandraProxyCreateDagForState: failed to unmarshal shared pool from dag args", zap.Error(err))
			continue
		}
		for _, component := range sharedPool.GetComponentGroups() {
			if component.GetId() != state.ClusterID || component.GetComponentType() != apiv1.ComponentType_CASSANDRA_PROXY {
				continue
			}
			if component.GetLabels()[cassandraProxyHostIdLabelKey] == state.HostID {
				return true
			}
		}
	}
	return false
}

func newCreateCassandraProxyDagExecID(ctx context.Context, sharedPoolID, clusterID, sharedPoolSnapshot, retryToken string) string {
	snapshotHash := utils.FnvHash(sharedPoolSnapshot)
	if retryToken != "" {
		snapshotHash = fmt.Sprintf("%s-%s", snapshotHash, utils.FnvHash(retryToken))
	}
	return utils.GeneratedDagExecId(ctx,
		fmt.Sprintf("create-cassandra-proxy-%s-%s-%s", sharedPoolID, clusterID, snapshotHash))
}

func buildCassandraProxySetResponse(state *entity.CassandraProxyState, phase apiv1.Component_Phase, image string) *apiv1.SetCassandraProxyResponse {
	return &apiv1.SetCassandraProxyResponse{
		HostId:             state.HostID,
		InternalEndpoint:   buildCassandraProxyInternalEndpoint(state.SharedPoolID, state.ClusterID),
		PrivatelinkDnsName: state.PrivatelinkDNSName,
		Phase:              phase,
		Image:              image,
	}
}

func (s *serverImpl) buildCassandraProxySetResponse(ctx context.Context, state *entity.CassandraProxyState, phase apiv1.Component_Phase, image string) *apiv1.SetCassandraProxyResponse {
	resp := buildCassandraProxySetResponse(state, phase, image)
	if resp.GetPrivatelinkDnsName() == "" && phase == apiv1.Component_AVAILABLE {
		resp.PrivatelinkDnsName = s.lookupCassandraProxyPLSDns(ctx, state.SharedPoolID, state.ClusterID)
	}
	return resp
}

func newCassandraProxyComponent(state *entity.CassandraProxyState, phase apiv1.Component_Phase) *entity.Component {
	component := &entity.Component{
		Id:            state.ClusterID,
		ComponentType: apiv1.ComponentType_CASSANDRA_PROXY,
		NodeQuota:     state.NodeQuotaEntity,
		Phase:         phase,
		Labels: map[string]string{
			cassandraProxyHostIdLabelKey:   state.HostID,
			cassandraProxyNodeSpecLabelKey: state.NodeSpec,
		},
	}
	return component
}

func cassandraProxySizingChanged(req *apiv1.SetCassandraProxyRequest, state *entity.CassandraProxyState, nodeSpec string) bool {
	return nodeSpec != state.NodeSpec ||
		(req.GetNodeQuota() != nil && !nodeQuotaEqual(state.NodeQuotaEntity, req.GetNodeQuota()))
}

func buildCassandraProxyInternalEndpoint(sharedPoolId, clusterId string) string {
	return fmt.Sprintf("%s.%s.svc.cluster.local:9042",
		entity.GetCassandraProxyResourceName(clusterId),
		entity.GetDataPlaneNamespace(sharedPoolId),
	)
}

func (s *serverImpl) handleCassandraProxyScale(
	ctx context.Context,
	req *apiv1.SetCassandraProxyRequest,
	state *entity.CassandraProxyState,
	nodeSpec string,
) (*apiv1.SetCassandraProxyResponse, error) {
	var (
		dagExecId          string
		dagSubmitted       bool
		capturedFreshState *entity.CassandraProxyState
		respImage          string
	)

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

		sharedPool, err := s.sharedPoolRepo.Load(ctx, state.SharedPoolID, false)
		if err != nil {
			if status.Code(err) == codes.NotFound {
				return status.Errorf(codes.NotFound, "shared pool %s not found", state.SharedPoolID)
			}
			return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
		}
		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 err := checkSharedPoolCanAcceptNewComponent(sharedPool); err != nil {
			return err
		}

		// TOCTOU recheck under lock: another goroutine may have changed state since outer load.
		freshState, err := s.cassandraProxyStateRepo.LoadByClusterID(ctx, req.GetClusterId(), false)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to reload cassandra-proxy state: %s", err)
		}
		if freshState == nil {
			return status.Errorf(codes.Internal, "cassandra-proxy state vanished for cluster %s", req.GetClusterId())
		}
		capturedFreshState = freshState

		freshPhase := entity.CassandraProxyPhaseToComponentPhase(freshState.Phase)
		newImage, freshImageChanged, err := s.resolveCassandraProxyImageForUpdate(req, freshState)
		if err != nil {
			return err
		}
		freshSizingChanged := cassandraProxySizingChanged(req, freshState, nodeSpec)
		switch freshPhase {
		case apiv1.Component_CREATING, apiv1.Component_MODIFYING:
			return status.Errorf(codes.FailedPrecondition,
				"cassandra proxy is %s, retry after AVAILABLE", freshPhase)
		case apiv1.Component_AVAILABLE:
			// Re-check diff: concurrent scale to same spec becomes a no-op.
			if !freshSizingChanged && !freshImageChanged {
				if newImage == "" {
					newImage, err = s.configuredCassandraProxyImage()
					if err != nil {
						return err
					}
				}
				respImage = newImage
				return nil
			}
		case apiv1.Component_UNAVAILABLE:
			// Preserve the same error code as the outer UNAVAILABLE path: changing
			// spec from a failed apply state is blocked until the current spec recovers.
			return status.Errorf(codes.FailedPrecondition, "cannot change cassandra-proxy desired state while component is UNAVAILABLE; retry with the current desired state first")
		default:
			return status.Errorf(codes.FailedPrecondition,
				"scale not supported in phase %s", freshPhase)
		}

		newNodeSpec := nodeSpec
		newNodeQuota := req.GetNodeQuota()
		if newImage == "" {
			newImage, err = s.configuredCassandraProxyImage()
			if err != nil {
				return err
			}
		}
		respImage = newImage

		newComp := newCassandraProxyComponent(&entity.CassandraProxyState{
			ClusterID:       freshState.ClusterID,
			HostID:          freshState.HostID,
			NodeSpec:        newNodeSpec,
			NodeQuotaEntity: newNodeQuota,
		}, apiv1.Component_MODIFYING)
		replaceComponentInGroups(sharedPool, newComp)
		sharedPool.Phase = apiv1.SharedPool_MODIFYING

		updatedState := &entity.CassandraProxyState{
			ClusterID:          freshState.ClusterID,
			SharedPoolID:       freshState.SharedPoolID,
			HostID:             freshState.HostID,
			NodeSpec:           newNodeSpec,
			Image:              newImage,
			Phase:              entity.ComponentPhaseToCassandraProxyPhase(apiv1.Component_MODIFYING),
			PrivatelinkDNSName: freshState.PrivatelinkDNSName,
			NodeQuotaEntity:    newNodeQuota,
		}
		if err := s.cassandraProxyStateRepo.Save(ctx, updatedState); err != nil {
			return status.Errorf(codes.Internal, "failed to save cassandra-proxy state: %s", err)
		}
		if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
			return status.Errorf(codes.Internal, "failed to save sharedPool: %s", err)
		}

		sharedPoolStr, err := protojson.Marshal(sharedPool)
		if err != nil {
			return 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(req.GetClusterId()),
			workflow.ComponentTypesKey: actions.ActionArgString(entity.GetComponentTypeName(newComp)),
			workflow.TiDBClusterIdKey:  actions.ActionArgString(req.GetClusterId()),
			utils.TraceIdKey:           utils.GetTraceId(ctx),
		}
		dagExecId = utils.GeneratedDagExecId(ctx,
			fmt.Sprintf("scale-cassandra-proxy-%s-%s-%s",
				sharedPool.GetId(), req.GetClusterId(), utils.FnvHash(string(sharedPoolStr))))

		if err := s.dagExecutor.Submit(ctx, dag.PlanCreateCassandraProxy, args, dagExecId); err != nil {
			return status.Errorf(codes.Internal, "failed to submit scale-cassandra-proxy dag: %s", err)
		}
		dagSubmitted = true
		tilog.Info(ctx, "submitted scale cassandra-proxy dag",
			zap.String("sharedPoolId", state.SharedPoolID),
			zap.String("clusterId", req.GetClusterId()),
			zap.String("dagExecId", dagExecId))
		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 nil, err
	}

	if !dagSubmitted {
		// No spec diff under lock (concurrent scale resolved it): return current AVAILABLE state.
		return s.buildCassandraProxySetResponse(ctx, capturedFreshState, apiv1.Component_AVAILABLE, respImage), nil
	}
	return &apiv1.SetCassandraProxyResponse{
		HostId:             capturedFreshState.HostID,
		InternalEndpoint:   buildCassandraProxyInternalEndpoint(capturedFreshState.SharedPoolID, req.GetClusterId()),
		PrivatelinkDnsName: capturedFreshState.PrivatelinkDNSName,
		Phase:              apiv1.Component_MODIFYING,
		Image:              respImage,
	}, nil
}

// replaceComponentInGroups replaces the CASSANDRA_PROXY component matching comp.GetId()
// in sp.ComponentGroups with comp. If no matching entry exists (state-only component),
// comp is appended so the snapshot remains consistent with durable state.
func replaceComponentInGroups(sp *entity.SharedPool, comp *entity.Component) {
	for i, c := range sp.GetComponentGroups() {
		if c.GetComponentType() == apiv1.ComponentType_CASSANDRA_PROXY &&
			c.GetId() == comp.GetId() {
			sp.ComponentGroups[i] = comp
			return
		}
	}
	sp.ComponentGroups = append(sp.ComponentGroups, comp)
}

// nodeQuotaEqual compares the scale-relevant fields of two NodeQuota values.
// Returns true when both quotas describe the same desired sizing.
func nodeQuotaEqual(a, b *apiv1.NodeQuota) bool {
	if a == nil || b == nil {
		return a == nil && b == nil
	}
	if a.GetNodeCount() != b.GetNodeCount() ||
		a.GetMinNodeCount() != b.GetMinNodeCount() ||
		a.GetMaxNodeCount() != b.GetMaxNodeCount() {
		return false
	}
	aSpec, bSpec := a.GetNodeSpec(), b.GetNodeSpec()
	if aSpec == nil || bSpec == nil {
		return aSpec == nil && bSpec == nil
	}
	return aSpec.GetCpu() == bSpec.GetCpu() &&
		aSpec.GetMemoryGi() == bSpec.GetMemoryGi()
}