package action

import (
	"context"
	"fmt"

	"github.com/google/uuid"
	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/pkg/dataplanemgr"
	"github.com/tidbcloud/cluster-service-ng/pkg/nodegroup"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"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"
)

const (
	// ReactivateNodeInstanceGroupActionName is the action name for reactivating node instance group.
	ReactivateNodeInstanceGroupActionName = "reactivate-node-instance-group-action"
)

// ReactivateNodeInstanceGroupAction reactivates a node instance group by removing deprecated taints.
type ReactivateNodeInstanceGroupAction struct {
	nodeGroupStateRepo     repo.DataplaneNodeInstanceGroupStateRepo
	nodeGroupLockRepo      repo.DataplaneNodeInstanceGroupLockRepo
	dataplaneServiceClient dplanemgrv1.DataPlaneServiceClient
}

// NewReactivateNodeInstanceGroupAction creates a new ReactivateNodeInstanceGroupAction.
func NewReactivateNodeInstanceGroupAction(
	nodeGroupStateRepo repo.DataplaneNodeInstanceGroupStateRepo,
	nodeGroupLockRepo repo.DataplaneNodeInstanceGroupLockRepo,
	dataplaneClient dataplanemgr.Client,
) actions.Action {
	return &ReactivateNodeInstanceGroupAction{
		nodeGroupStateRepo:     nodeGroupStateRepo,
		nodeGroupLockRepo:      nodeGroupLockRepo,
		dataplaneServiceClient: dataplaneClient.GetDataPlaneServiceClient(),
	}
}

// Name returns the action name.
func (a *ReactivateNodeInstanceGroupAction) Name() string {
	return ReactivateNodeInstanceGroupActionName
}

// Invoke executes the reactivate node instance group action.
func (a *ReactivateNodeInstanceGroupAction) Invoke(args actions.ActionArgs) error {
	ctx := utils.GenerateContext(args)
	tilog.Info(ctx, "reactivate node instance group action invoked")

	dataplaneName, err := args.GetString(workflow.DataplaneIdKey)
	if err != nil {
		return fmt.Errorf("failed to get dataplane_id: %w", err)
	}

	groupData, err := args.GetString(workflow.NodeInstanceGroupKey)
	if err != nil {
		return fmt.Errorf("failed to get node_instance_group: %w", err)
	}
	var group entity.NodeInstanceGroup
	if err := protojson.Unmarshal([]byte(groupData), &group); err != nil {
		return status.Errorf(codes.InvalidArgument, "failed to unmarshal node instance group: %v", err)
	}

	nodeSpec := entity.GetComponentNodeSpec(group.GetNodeSpec())
	primaryInstance := group.GetPrimaryInstanceType()

	if dataplaneName == "" {
		return status.Errorf(codes.InvalidArgument, "dataplane_id is required")
	}
	if err := nodegroup.ValidateSpec(nodeSpec); err != nil {
		return status.Errorf(codes.InvalidArgument, "invalid node_spec: %v", err)
	}
	if err := nodegroup.ValidateCloudInstanceType(primaryInstance); err != nil {
		return status.Errorf(codes.InvalidArgument, "invalid primary_instance: %v", err)
	}

	ownerID := fmt.Sprintf("%s-%s-%s", ReactivateNodeInstanceGroupActionName, nodeSpec, uuid.New().String()[:8])
	lock, err := a.nodeGroupLockRepo.TryAcquire(ctx, dataplaneName, nodeSpec, primaryInstance, ownerID, DefaultLockLeaseDuration)
	if err != nil {
		tilog.Error(ctx, "failed to acquire lock",
			zap.String("dataplaneName", dataplaneName),
			zap.String("nodeSpec", nodeSpec),
			zap.String("ownerID", ownerID),
			zap.Error(err))
		return status.Errorf(codes.Unavailable, "failed to acquire lock: %v", err)
	}
	if lock == nil {
		tilog.Warn(ctx, "lock is held by another owner",
			zap.String("dataplaneName", dataplaneName),
			zap.String("nodeSpec", nodeSpec),
			zap.String("ownerID", ownerID))
		return status.Errorf(codes.Aborted, "lock is held by another owner, retry later")
	}

	defer func() {
		if releaseErr := a.nodeGroupLockRepo.Release(ctx, dataplaneName, nodeSpec, primaryInstance, ownerID); releaseErr != nil {
			tilog.Warn(ctx, "failed to release lock",
				zap.String("dataplaneName", dataplaneName),
				zap.String("nodeSpec", nodeSpec),
				zap.String("ownerID", ownerID),
				zap.Error(releaseErr))
		}
	}()

	return a.reactivateNodeInstanceGroup(ctx, dataplaneName, nodeSpec, primaryInstance, group.GetFallbackInstances())
}

func (a *ReactivateNodeInstanceGroupAction) reactivateNodeInstanceGroup(
	ctx context.Context,
	dataplaneName, nodeSpec, primaryInstance string,
	fallbackInstances []*entity.NodeInstanceGroupFallbackInstance,
) error {
	nodePoolsResp, err := a.dataplaneServiceClient.ListNodePools(ctx, &dplanemgrv1.ListNodePoolsRequest{
		DataPlaneId: dataplaneName,
	})
	if err != nil {
		return status.Errorf(codes.Internal, "failed to list nodepools: %v", err)
	}

	groupInstanceTypes := map[string]nodegroup.NodePoolRole{
		primaryInstance: nodegroup.NodePoolRolePrimary,
	}
	for _, fb := range fallbackInstances {
		if fb != nil && fb.GetInstanceType() != "" {
			if _, exists := groupInstanceTypes[fb.GetInstanceType()]; exists {
				continue
			}
			groupInstanceTypes[fb.GetInstanceType()] = nodegroup.NodePoolRoleFallback
		}
	}

	cpu, mem, err := nodegroup.ParseSpec(nodeSpec)
	if err != nil {
		return status.Errorf(codes.InvalidArgument, "invalid node spec %s: %v", nodeSpec, err)
	}

	deprecatedTaint := nodegroup.DeprecatedTaint()
	for _, np := range nodePoolsResp.GetNodePools() {
		if !matchesNodeInstanceGroupNodePool(np, nodeSpec, cpu, mem) {
			continue
		}
		expectedRole, ok := groupInstanceTypes[np.GetInstanceType()]
		if !ok {
			continue
		}
		if !nodegroup.NodePoolMatchesRole(np, expectedRole) {
			continue
		}
		if err := a.removeDeprecatedTaint(ctx, dataplaneName, np, deprecatedTaint); err != nil {
			return err
		}
	}

	state, err := a.nodeGroupStateRepo.LoadOrCreate(ctx, dataplaneName, nodeSpec, primaryInstance)
	if err != nil {
		return status.Errorf(codes.Internal, "failed to load or create state: %v", err)
	}
	if state.State == entity.NodeInstanceGroupStateDeprecated {
		if err := a.nodeGroupStateRepo.TransitionState(ctx, state.ID, entity.NodeInstanceGroupStateDeprecated, entity.NodeInstanceGroupStateNormal); err != nil {
			return status.Errorf(codes.Internal, "failed to transition state to normal: %v", err)
		}
	}

	return nil
}

func (a *ReactivateNodeInstanceGroupAction) removeDeprecatedTaint(
	ctx context.Context,
	dataplaneName string,
	nodePool *dplanemgrv1.NodePool,
	taint nodegroup.Taint,
) error {
	existingTaints := nodePoolTaintsToNodeGroup(nodePool.GetTaints())
	if !nodegroup.HasTaint(existingTaints, taint.Key, taint.Value) {
		return nil
	}

	updatedTaints := nodegroup.RemoveTaintByKeyValue(existingTaints, taint.Key, taint.Value)
	req := &dplanemgrv1.UpdateNodePoolRequest{
		DataPlaneId: dataplaneName,
		NodePoolId:  nodePool.GetName(),
		NodePool: &dplanemgrv1.NodePool{
			// Preserve existing labels because UpdateNodePool has replace semantics.
			Labels: nodePool.GetLabels(),
			Taints: nodeGroupTaintsToNodePool(updatedTaints),
		},
	}
	if _, err := a.dataplaneServiceClient.UpdateNodePool(ctx, req); err != nil {
		tilog.Error(ctx, "failed to remove deprecated taint from nodepool",
			zap.String("dataplaneName", dataplaneName),
			zap.String("nodePoolName", nodePool.GetName()),
			zap.String("taintKey", taint.Key),
			zap.String("taintValue", taint.Value),
			zap.Error(err))
		return status.Errorf(codes.Internal, "failed to update nodepool %s taint: %v", nodePool.GetName(), err)
	}

	tilog.Info(ctx, "removed deprecated taint from nodepool",
		zap.String("dataplaneName", dataplaneName),
		zap.String("nodePoolName", nodePool.GetName()),
		zap.String("taintKey", taint.Key),
		zap.String("taintValue", taint.Value))
	return nil
}

func matchesNodeInstanceGroupNodePool(nodePool *dplanemgrv1.NodePool, nodeSpec string, cpuCore, memoryGi uint32) bool {
	if nodePool == nil {
		return false
	}

	labels := nodePool.GetLabels()
	if nodegroup.IsManagedByInstanceTypeManager(labels) {
		return labels != nil && labels[nodegroup.LabelNodeSpec] == nodeSpec
	}
	return nodePool.GetCpuCore() == cpuCore && nodePool.GetMemoryGi() == memoryGi
}