package server

import (
	"context"
	"fmt"
	"strings"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"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) SetSharedPoolComponentNodeVersion(ctx context.Context, sharedPoolId, compId string, compType apiv1.ComponentType, nodeVersion string) (*entity.Component, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "sharedPool id is required")
	}

	if compType == apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
		return nil, status.Errorf(codes.InvalidArgument, "component type is required")
	}

	if nodeVersion == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_version is required")
	}

	var currComponent *entity.Component
	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 pool 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)
			}
			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)
		}

		if sharedPool.GetPhase() == apiv1.SharedPool_CREATING {
			return status.Errorf(codes.FailedPrecondition, "shared pool is creating, please try again later")
		}

		if compId == "" {
			tilog.Warn(ctx, "target component id is empty, will use default id")
			compId = sharedPool.GetDefaultComponenId()
		}

		// Find the target component
		var found bool
		for _, comp := range sharedPool.GetComponentGroups() {
			if comp.GetId() != compId || comp.GetComponentType() != compType {
				continue
			}

			found = true
			currComponent = comp
			break
		}

		if !found {
			return status.Errorf(codes.NotFound, "component with id '%s' and type '%s' not found in sharedPool",
				compId, compType.String())
		}

		if err := checkComponentDeleteFence(currComponent); err != nil {
			return err
		}

		if isTiKVRemoteWorkerGuardedComponentType(compType) {
			if err := s.guardTiKVRemoteWorkersTopologyStable(ctx, sharedPool); err != nil {
				return err
			}
		}
		currComponent.Phase = apiv1.Component_MODIFYING

		// Validate nodeVersion exists in node pools for the component's node spec
		if err := s.validateNodeVersionForComponent(ctx, sharedPool.GetDataplaneName(), currComponent, nodeVersion); err != nil {
			return err
		}

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

		// Build target component for workflow
		targetComponent := &entity.Component{
			Id:            compId,
			ComponentType: compType,
		}

		// start update component dag
		args := actions.ActionArgs{
			workflow.SharedPoolIdKey:                      actions.ActionArgString(sharedPool.GetId()),
			workflow.SharedPoolKey:                        actions.ActionArgString(string(sharedPoolStr)),
			workflow.ComponentIdsKey:                      actions.ActionArgString(compId),
			workflow.ComponentTypesKey:                    actions.ActionArgString(entity.GetComponentTypeName(targetComponent)),
			entity.GetUpdateComponentKey(targetComponent): actions.ActionArgString(string(tioperator.OpsUpdateNodeVersion)),
			workflow.NodeVersionKey:                       actions.ActionArgString(nodeVersion),
			utils.TraceIdKey:                              utils.GetTraceId(ctx),
		}

		dagExecId = utils.GeneratedDagExecId(ctx,
			fmt.Sprintf("set-component-node-version-%s-%s-%s-%s",
				sharedPool.GetId(),
				entity.GetComponentTypeName(targetComponent),
				compId,
				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-node-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 nil, err
	}

	return currComponent, nil
}

// validateNodeVersionForComponent validates that the given nodeVersion exists in the node pools
// for the component's node spec (cpu, memory, and local disk configuration).
func (s *serverImpl) validateNodeVersionForComponent(ctx context.Context, dataplaneName string, comp *entity.Component, nodeVersion string) error {
	// Validate component has valid node quota and node spec
	if comp.GetNodeQuota() == nil || comp.GetNodeQuota().GetNodeSpec() == nil {
		return status.Errorf(codes.Internal, "component has invalid node quota")
	}

	// Get node pools from dataplane manager
	nodePoolResp, err := s.dataplaneClient.ListNodePools(ctx, &dplanemgrv1.ListNodePoolsRequest{
		DataPlaneId: dataplaneName,
	})
	if err != nil {
		tilog.Error(ctx, "failed to list node pools", zap.String("dataplaneName", dataplaneName), zap.Error(err))
		return status.Errorf(codes.Internal, "failed to list node pools: %s", err)
	}

	// Build node spec key for the component
	nodeSpec := entity.GetComponentNodeSpec(comp.GetNodeQuota().GetNodeSpec())
	for _, storage := range comp.GetStorages() {
		if storage.GetName() != apiv1.Component_Storage_DATA {
			continue
		}
		if storage.GetType() == entity.LocalSSD {
			nodeSpec = fmt.Sprintf("%s-%s", nodeSpec, "Local")
		}
		break
	}

	// Check if nodeVersion exists for the component's node spec
	var availableVersions []string
	for _, np := range nodePoolResp.GetNodePools() {
		if np.GetNodeType() != dplanemgrv1.NodeType_Dedicated {
			continue
		}

		key := fmt.Sprintf("%dC%dG", np.GetCpuCore(), np.GetMemoryGi())
		if np.GetWithLocalDisk() {
			key = fmt.Sprintf("%s-%s", key, "Local")
		}

		if key != nodeSpec {
			continue
		}

		// Found matching node spec, check if nodeVersion matches
		if np.GetK8SVersion() == nodeVersion {
			return nil // nodeVersion is valid
		}
		availableVersions = append(availableVersions, np.GetK8SVersion())
	}

	return status.Errorf(codes.InvalidArgument, "node version '%s' not available for node spec '%s', available versions: [%s]", nodeVersion, nodeSpec, strings.Join(availableVersions, ", "))
}