package server

import (
	"context"
	"fmt"
	"net"
	"reflect"
	"strconv"
	"strings"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	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) UpdateTiDBGroup(ctx context.Context, tidbGroup *entity.TiDBGroup, force bool) (*entity.TiDBGroup, error) {
	if err := s.validateUpdateTiDBGroupRequest(tidbGroup); err != nil {
		return nil, err
	}

	var currTiDBGroup *entity.TiDBGroup
	var dagExecId string
	var dagSubmitted bool

	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		var err error
		currTiDBGroup, err = s.prepareTiDBGroupUpdate(ctx, tidbGroup)
		if err != nil {
			return err
		}

		sharedPool, unlockFunc, err := s.lockAndLoadSharedPool(ctx, currTiDBGroup.GetSharedPoolId())
		if err != nil {
			return err
		}
		defer func() {
			if unlockErr := unlockFunc(); unlockErr != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", currTiDBGroup.GetSharedPoolId()), zap.Error(unlockErr))
			}
		}()

		if err := s.checkTiDBGroupConcurrency(ctx, currTiDBGroup.GetSharedPoolId(), tidbGroup.GetId(), force); err != nil {
			return err
		}

		currTiDBComp, err := s.findTiDBComponent(sharedPool, currTiDBGroup.GetId())
		if err != nil {
			return err
		}

		if currTiDBComp.GetPhase() == apiv1.Component_CREATING {
			return status.Errorf(codes.FailedPrecondition, "tidb group is creating, can not update")
		}

		optionsMap := map[string][]string{
			entity.GetUpdateComponentKey(currTiDBComp): {},
		}
		updateLb := false

		if tidbGroup.GetSqlEndpoint() != nil {
			lbUpdated, err := s.updateTiDBGroupEndpoint(ctx, tidbGroup, currTiDBGroup, sharedPool, currTiDBComp, optionsMap)
			if err != nil {
				return err
			}
			updateLb = lbUpdated
		}

		if tidbGroup.GetNodeQuota() != nil {
			if tidbGroup.GetNodeQuota().GetNodeCount() != 0 && tidbGroup.GetNodeQuota().GetNodeCount() != currTiDBComp.GetNodeQuota().GetNodeCount() {
				if err := s.checkManualScaleHPAGate(
					ctx,
					sharedPool,
					apiv1.ComponentType_TIDB,
					tidbGroup.GetId(),
					currTiDBComp.GetNodeQuota().GetNodeCount(),
					tidbGroup.GetNodeQuota().GetNodeCount(),
					force,
				); err != nil {
					return err
				}
			}
			if err := s.updateTiDBGroupNodeQuota(ctx, tidbGroup, currTiDBGroup, sharedPool, currTiDBComp, optionsMap, force); err != nil {
				return err
			}
		}

		currTiDBGroup.Phase = apiv1.TiDBGroup_MODIFYING

		dagExecId, dagSubmitted, err = s.submitUpdateTiDBGroupDAG(ctx, currTiDBGroup, sharedPool, currTiDBComp, updateLb, 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 currTiDBGroup, nil
}

func (s *serverImpl) validateUpdateTiDBGroupRequest(tidbGroup *entity.TiDBGroup) error {
	if tidbGroup == nil || tidbGroup.GetId() == "" {
		return status.Errorf(codes.InvalidArgument, "tidb group id is empty")
	}
	return nil
}

func (s *serverImpl) prepareTiDBGroupUpdate(ctx context.Context, tidbGroup *entity.TiDBGroup) (*entity.TiDBGroup, error) {
	currTiDBGroup, err := s.GetTiDBGroup(ctx, tidbGroup.GetClusterId(), tidbGroup.GetId())
	if err != nil {
		return nil, err
	}

	tidbGroup.SharedPoolId = currTiDBGroup.GetSharedPoolId()
	tidbGroup.ClusterId = currTiDBGroup.GetClusterId()
	tidbGroup.TenantId = currTiDBGroup.GetTenantId()

	return currTiDBGroup, nil
}

func (s *serverImpl) lockAndLoadSharedPool(ctx context.Context, sharedPoolId string) (*entity.SharedPool, func() error, 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 nil, nil, err
	}

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

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

	return sharedPool, unlockFunc, nil
}

func (s *serverImpl) checkTiDBGroupConcurrency(ctx context.Context, sharedPoolId, tidbGroupId string, force bool) error {
	dagIns, err := s.getRunningDagIns(ctx, []string{
		dag.PlanCreateTiDBGroup,
		dag.PlanUpdateTiDBGroup,
		dag.PlanDeleteTiDBGroup,
	}, []string{sharedPoolId}, nil, nil, []string{tidbGroupId})
	if err != nil {
		return status.Errorf(codes.Internal, "failed to get running tidb group dag instance: %s", err)
	}

	for _, dagIns := range dagIns {
		switch dagIns.ExecTask.RootName {
		case dag.PlanCreateTiDBGroup:
			return status.Errorf(codes.FailedPrecondition, "tidb group is creating, please try again later")
		case dag.PlanUpdateTiDBGroup:
			if !force {
				return status.Errorf(codes.FailedPrecondition, "tidb group is updating, please try again later")
			}
		case dag.PlanDeleteTiDBGroup:
			return status.Errorf(codes.FailedPrecondition, "tidb group is deleting")
		}
	}
	return nil
}

func (s *serverImpl) findTiDBComponent(sharedPool *entity.SharedPool, tidbGroupId string) (*entity.Component, error) {
	for _, comp := range sharedPool.GetComponentGroups() {
		if comp.GetComponentType() == apiv1.ComponentType_TIDB && comp.GetId() == tidbGroupId {
			return comp, nil
		}
	}
	return nil, status.Errorf(codes.NotFound, "tidb group not found in sharedPool")
}

func (s *serverImpl) updateTiDBGroupEndpoint(ctx context.Context, tidbGroup, currTiDBGroup *entity.TiDBGroup, sharedPool *entity.SharedPool, currTiDBComp *entity.Component, optionsMap map[string][]string) (bool, error) {
	if err := checkTidbGroupEndpoint(tidbGroup); err != nil {
		return false, err
	}

	if tidbGroup.GetSqlEndpoint().GetPrivateLb() != nil {
		currTiDBGroup.GetSqlEndpoint().PrivateLb = tidbGroup.GetSqlEndpoint().GetPrivateLb()
	}

	if tidbGroup.GetSqlEndpoint().GetPublicLb() != nil {
		currTiDBGroup.GetSqlEndpoint().PublicLb = tidbGroup.GetSqlEndpoint().GetPublicLb()

		tiproxyComp := s.findTiProxyComponent(sharedPool, currTiDBGroup.GetClusterId())
		if tiproxyComp == nil {
			tilog.Error(ctx, "tiproxy component not found in shared pool", zap.String("sharedPoolId", sharedPool.GetId()), zap.String("clusterId", currTiDBGroup.GetClusterId()))
			return false, status.Errorf(codes.Internal, "tiproxy component not found in shared pool")
		}

		tiproxyComp.Phase = apiv1.Component_MODIFYING
		optionsMap[entity.GetUpdateComponentKey(tiproxyComp)] = []string{string(tioperator.OpsUpdateConfig), string(tioperator.OpsUpdateUpdateStrategy)}
	}

	currTiDBComp.Phase = apiv1.Component_MODIFYING
	optionsMap[entity.GetUpdateComponentKey(currTiDBComp)] = append(optionsMap[entity.GetUpdateComponentKey(currTiDBComp)],
		string(tioperator.OpsUpdateLabelCidr))

	return true, nil
}

func (s *serverImpl) findTiProxyComponent(sharedPool *entity.SharedPool, clusterId string) *entity.Component {
	return entity.FindTiProxyComponent(sharedPool, clusterId)
}

func (s *serverImpl) updateTiDBGroupNodeQuota(ctx context.Context, tidbGroup, currTiDBGroup *entity.TiDBGroup, sharedPool *entity.SharedPool, currTiDBComp *entity.Component, optionsMap map[string][]string, force bool) error {
	// Save original values before any mutation for scale-in safety check
	origNodeCount := currTiDBComp.GetNodeQuota().GetNodeCount()

	origNodeSpec := currTiDBComp.GetNodeQuota().GetNodeSpec()

	// Determine target node count
	// Note: proto zero-value (0) is treated as "not provided" for spec-only update requests.
	// This is safe because TiDB does not support scaling to zero nodes (min node count >= 1).
	// For TiFlash, zero node count is valid but handled via explicit field check, not zero-value.
	targetNodeCount := tidbGroup.GetNodeQuota().GetNodeCount()
	if targetNodeCount == 0 {
		targetNodeCount = origNodeCount
	}
	targetNodeMin := tidbGroup.GetNodeQuota().GetMinNodeCount()
	if targetNodeMin == 0 {
		targetNodeMin = currTiDBComp.GetNodeQuota().GetMinNodeCount()
	}
	targetNodeMax := tidbGroup.GetNodeQuota().GetMaxNodeCount()
	if targetNodeMax == 0 {
		targetNodeMax = currTiDBComp.GetNodeQuota().GetMaxNodeCount()
	}

	hasCountChange := tidbGroup.GetNodeQuota().GetNodeCount() != 0 && targetNodeCount != origNodeCount

	// Step 1: Deterministic validations first (local, cheap, fail fast)
	if hasCountChange {
		if err := s.validateNodeCountUpdate(ctx, tidbGroup, currTiDBGroup, currTiDBComp, force); err != nil {
			return err
		}
	}

	// Step 2: Resolve target node spec (if spec change requested)
	var targetNodeSpec *entity.NodeSpec
	hasSpecChange := tidbGroup.GetNodeQuota().GetNodeSpec() != nil

	if hasSpecChange {
		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)
		}

		resolvedSpec, err := getNodeSpec(ctx, tidbGroup.GetNodeQuota().GetNodeSpec(), nodeSpecMap, false)
		if err != nil {
			tilog.Error(ctx, "failed to get new node spec", zap.Error(err))
			return status.Errorf(codes.InvalidArgument, "failed to get new node spec: %s", err)
		}
		targetNodeSpec = resolvedSpec
	} else {
		targetNodeSpec = origNodeSpec
	}

	// Step 3: Run RU validation (local, cheap) before remote metrics fetch
	if hasSpecChange {
		ruConfig, err := s.ruConfigRepo.GetRuConfig(ctx, sharedPool.GetId(), apiv1.RuConfig_SHARED_POOL)
		if err != nil {
			tilog.Error(ctx, "failed to get ru config", zap.String("sharedPoolId", sharedPool.GetId()), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to get ru config: %s", err)
		}

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

		if err := validateNodeSpecAgainstRUConfig(ctx, currTiDBComp.GetComponentType(), targetNodeSpec, ruConfig, minRU, currTiDBComp.GetId(), force); err != nil {
			return err
		}
	}

	// Step 4: Run checkScaleInSafety exactly once with resolved target spec
	// This prevents false rejection on count-down + spec-up (e.g., 3x4c -> 2x16c)
	if err := s.checkScaleInSafety(ctx, sharedPool, currTiDBComp,
		origNodeCount, origNodeSpec,
		targetNodeCount, targetNodeSpec, force); err != nil {
		return err
	}

	// Step 5: Mutate state
	if hasCountChange {
		currTiDBComp.GetNodeQuota().NodeCount = targetNodeCount
		currTiDBComp.Phase = apiv1.Component_MODIFYING
		if currTiDBComp.GetNodeQuota().MaxNodeCount != targetNodeMax {
			currTiDBComp.GetNodeQuota().MaxNodeCount = targetNodeMax
			optionsMap[entity.GetUpdateComponentKey(currTiDBComp)] = append(optionsMap[entity.GetUpdateComponentKey(currTiDBComp)],
				string(tioperator.OpsUpdateMaxCount))
		}
		if currTiDBComp.GetNodeQuota().MinNodeCount != targetNodeMin {
			currTiDBComp.GetNodeQuota().MinNodeCount = targetNodeMin
			optionsMap[entity.GetUpdateComponentKey(currTiDBComp)] = append(optionsMap[entity.GetUpdateComponentKey(currTiDBComp)],
				string(tioperator.OpsUpdateMinCount))
		}
		optionsMap[entity.GetUpdateComponentKey(currTiDBComp)] = append(optionsMap[entity.GetUpdateComponentKey(currTiDBComp)],
			string(tioperator.OpsUpdateNodeCount))
	}

	if hasSpecChange && !reflect.DeepEqual(currTiDBComp.GetNodeQuota().GetNodeSpec(), targetNodeSpec) {
		tilog.Info(ctx, "update tidb group node spec",
			zap.String("oldNodeSpec", currTiDBComp.GetNodeQuota().GetNodeSpec().String()),
			zap.String("newNodeSpec", targetNodeSpec.String()),
		)

		if currTiDBGroup.GetPhase() != apiv1.TiDBGroup_AVAILABLE {
			if !force {
				tilog.Warn(ctx, "tidb group is not available, can not update node spec",
					zap.String("tidbGroupId", currTiDBGroup.GetId()),
					zap.String("phase", currTiDBGroup.GetPhase().String()),
				)
				return status.Errorf(codes.FailedPrecondition, "tidb group phase is %s, can not update node spec",
					currTiDBGroup.GetPhase().String())
			}
			tilog.Warn(ctx, "tidb group is not available, force update node spec",
				zap.String("tidbGroupId", currTiDBGroup.GetId()),
				zap.String("phase", currTiDBGroup.GetPhase().String()),
			)
		}
		currTiDBComp.GetNodeQuota().NodeSpec = targetNodeSpec
		currTiDBComp.Phase = apiv1.Component_MODIFYING
		optionsMap[entity.GetUpdateComponentKey(currTiDBComp)] = append(optionsMap[entity.GetUpdateComponentKey(currTiDBComp)],
			string(tioperator.OpsUpdateNodeSpec))
	}

	currTiDBGroup.NodeQuota = currTiDBComp.GetNodeQuota()
	return nil
}

func (s *serverImpl) validateNodeCountUpdate(ctx context.Context, tidbGroup, currTiDBGroup *entity.TiDBGroup, currTiDBComp *entity.Component, force bool) error {
	if currTiDBGroup.GetPhase() != apiv1.TiDBGroup_AVAILABLE {
		if !force {
			return status.Errorf(codes.FailedPrecondition, "tidb group phase is %s, can not update node count", currTiDBGroup.GetPhase().String())
		}
		tilog.Warn(ctx, "tidb group is not available, force update node count",
			zap.String("tidbGroupId", currTiDBGroup.GetId()),
			zap.String("phase", currTiDBGroup.GetPhase().String()),
		)
	}

	if tidbGroup.GetNodeQuota().GetMaxNodeCount() > 0 &&
		tidbGroup.GetNodeQuota().GetNodeCount() > tidbGroup.GetNodeQuota().GetMaxNodeCount() {
		return status.Errorf(codes.InvalidArgument, "tidb group node count is larger than new max node count")
	}
	if tidbGroup.GetNodeQuota().GetMaxNodeCount() == 0 &&
		currTiDBComp.GetNodeQuota().GetMaxNodeCount() < tidbGroup.GetNodeQuota().GetNodeCount() {
		return status.Errorf(codes.InvalidArgument, "tidb group node count is larger than current max node count")
	}

	if tidbGroup.GetNodeQuota().GetMinNodeCount() > 0 &&
		tidbGroup.GetNodeQuota().GetNodeCount() < tidbGroup.GetNodeQuota().GetMinNodeCount() {
		return status.Errorf(codes.InvalidArgument, "tidb group node count is less than new min node count")
	}

	if tidbGroup.GetNodeQuota().GetMinNodeCount() == 0 &&
		currTiDBComp.GetNodeQuota().GetMinNodeCount() > tidbGroup.GetNodeQuota().GetNodeCount() {
		return status.Errorf(codes.InvalidArgument, "tidb group node count is less than current min node count")
	}

	return nil
}

func (s *serverImpl) submitUpdateTiDBGroupDAG(ctx context.Context, currTiDBGroup *entity.TiDBGroup, sharedPool *entity.SharedPool, currTiDBComp *entity.Component, updateLb bool, 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 shard pool: %s", err)
	}

	tidbGroupStr, err := protojson.Marshal(currTiDBGroup)
	if err != nil {
		return "", false, status.Errorf(codes.Internal, "failed to marshal tidb group: %s", err)
	}

	if err := s.tidbGroupRepo.Save(ctx, currTiDBGroup); err != nil {
		tilog.Error(ctx, "failed to save tidb group meta", zap.Error(err))
		return "", false, fmt.Errorf("failed to save tidb group meta: %w", err)
	}

	componentIds := []string{currTiDBGroup.GetId()}
	componentTypes := []string{entity.GetComponentTypeName(currTiDBComp)}
	if updateLb {
		componentIds = append(componentIds, currTiDBGroup.GetClusterId())
		componentTypes = append(componentTypes, entity.GetComponentTypeName(&entity.Component{ComponentType: apiv1.ComponentType_TIPROXY}))
	}

	args := actions.ActionArgs{
		workflow.SharedPoolIdKey:       actions.ActionArgString(currTiDBGroup.GetSharedPoolId()),
		workflow.SharedPoolKey:         actions.ActionArgString(string(sharedPoolStr)),
		workflow.ComponentIdsKey:       actions.ActionArgString(strings.Join(componentIds, ",")),
		workflow.ComponentTypesKey:     actions.ActionArgString(strings.Join(componentTypes, ",")),
		workflow.TenantIdKey:           actions.ActionArgString(currTiDBGroup.GetTenantId()),
		workflow.TiDBGroupIdKey:        actions.ActionArgString(currTiDBGroup.GetId()),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
		workflow.TiDBClusterIdKey:      actions.ActionArgString(currTiDBGroup.GetClusterId()),
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(strconv.FormatBool(updateLb)),
		workflow.ForceKey:              actions.ActionArgString(strconv.FormatBool(force)),
		workflow.UpdateStrategyKey:     actions.ActionArgString(string(v1alpha1.ConfigUpdateStrategyHotReload)),
		utils.TraceIdKey:               utils.GetTraceId(ctx),
	}

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

	dagExecId := utils.GeneratedDagExecId(ctx, fmt.Sprintf("update-tidbgroup-%s-%s-%s",
		currTiDBGroup.GetSharedPoolId(),
		currTiDBGroup.GetId(),
		utils.FnvHash(string(tidbGroupStr)),
	))

	if err := s.dagExecutor.Submit(ctx, dag.PlanUpdateTiDBGroup, args, dagExecId); err != nil {
		return "", false, status.Errorf(codes.Internal, "failed to submit update tidb group dag: %s", err)
	}

	return dagExecId, true, nil
}

func checkTidbGroupEndpoint(tidbGroup *entity.TiDBGroup) error {
	if tidbGroup.GetSqlEndpoint().GetPrivateLb() != nil {
		if tidbGroup.GetSqlEndpoint().GetPrivateLb().GetPort() == 0 && tidbGroup.GetSqlEndpoint().GetPrivateLb().GetEnabled() {
			return status.Errorf(codes.InvalidArgument, "tidb group private lb port is required when enabled")
		}

		if tidbGroup.GetSqlEndpoint().GetPrivateLb().GetEnabled() {
			if err := validateIpAllowList(tidbGroup.GetSqlEndpoint().GetPrivateLb().GetIpAllowList(), "private"); err != nil {
				return err
			}
		}
	}

	if tidbGroup.GetSqlEndpoint().GetPublicLb() != nil {
		if tidbGroup.GetSqlEndpoint().GetPublicLb().GetPort() == 0 && tidbGroup.GetSqlEndpoint().GetPublicLb().GetEnabled() {
			return status.Errorf(codes.InvalidArgument, "tidb group public lb port is required when enabled")
		}

		if tidbGroup.GetSqlEndpoint().GetPublicLb().GetEnabled() {
			if err := validateIpAllowList(tidbGroup.GetSqlEndpoint().GetPublicLb().GetIpAllowList(), "public"); err != nil {
				return err
			}
		}
	}
	return nil
}

func validateIpAllowList(ipAllowLists []*entity.IpAllowList, lbType string) error {
	for _, ipAllowList := range ipAllowLists {
		cidr := strings.TrimSpace(ipAllowList.GetCidrBlock())
		if cidr == "" {
			return status.Errorf(codes.InvalidArgument, "tidb group %s lb ip allow list cidr block is required when enabled", lbType)
		}
		_, _, err := net.ParseCIDR(cidr)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "tidb group %s lb ip allow list cidr block %s is invalid: %s", lbType, cidr, err)
		}
		ipAllowList.CidrBlock = cidr
	}

	return nil
}