package action

import (
	"context"
	"fmt"
	"path"
	"strings"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	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/config"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/pkg/clusterconfig"
	"github.com/tidbcloud/cluster-service-ng/pkg/dataplane"
	"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/tioperator"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type UpdateComponentGroupAction struct {
	dataplaneManager    dataplane.Manager
	sharedPoolRepo      repo.SharedPoolRepo
	tidbGroupRepo       repo.TiDBGroupRepo
	tidbClusterRepo     repo.TiDBClusterRepo
	lbClient            dplanemgrv1.LoadBalancerServiceClient
	dataplaneClient     dplanemgrv1.DataPlaneServiceClient
	imageConfig         *config.ImageConfig
	clusterConfigClient clusterconfig.Client
	meteringConfig      *config.MeteringConfig
}

const UpdateComponentGroupActionName = "update-component-group"

func NewUpdateComponentGroupAction(
	sharedPoolRepo repo.SharedPoolRepo,
	tidbGroupRepo repo.TiDBGroupRepo,
	tidbClusterRepo repo.TiDBClusterRepo,
	dataplaneManager dataplane.Manager,
	clusterConfigClient clusterconfig.Client,
	lbClient dplanemgrv1.LoadBalancerServiceClient,
	dataplaneClient dplanemgrv1.DataPlaneServiceClient,
	imageConfig *config.ImageConfig,
	meteringConfig *config.MeteringConfig,
) actions.Action {
	return &UpdateComponentGroupAction{
		sharedPoolRepo:      sharedPoolRepo,
		tidbGroupRepo:       tidbGroupRepo,
		tidbClusterRepo:     tidbClusterRepo,
		dataplaneManager:    dataplaneManager,
		dataplaneClient:     dataplaneClient,
		clusterConfigClient: clusterConfigClient,
		lbClient:            lbClient,
		imageConfig:         imageConfig,
		meteringConfig:      meteringConfig,
	}
}

func (a *UpdateComponentGroupAction) Name() string {
	return UpdateComponentGroupActionName
}

func (a *UpdateComponentGroupAction) Invoke(args actions.ActionArgs) error {
	// generate context
	ctx := utils.GenerateContext(args)
	tilog.Info(ctx, "update component group action")

	_, err := args.GetString(workflow.ComponentIdsKey)
	if err != nil {
		return fmt.Errorf("target component id is required, but not found")
	}

	sharedPoolId, err := args.GetString(workflow.SharedPoolIdKey)
	if err != nil {
		return fmt.Errorf("failed to get sharedPool id: %w", err)
	}

	// load sharedPool
	sharedPool, err := a.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return fmt.Errorf("failed to load sharedPool: %w", err)
	}

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

	var tempSharedPool entity.SharedPool
	err = protojson.Unmarshal([]byte(sharedPoolStr), &tempSharedPool)
	if err != nil {
		tilog.Error(ctx, "failed to unmarshal sharedPool", zap.Error(err))
		return fmt.Errorf("failed to unmarshal sharedPool: %w", err)
	}

	sharedPool.ComponentGroups = tempSharedPool.GetComponentGroups()
	if sharedPool.GetVersion() == "" {
		sharedPool.Version = tempSharedPool.GetVersion()
	}

	remoteCli, err := a.dataplaneManager.GetClient(sharedPool.GetDataplaneName())
	if err != nil {
		tilog.Error(ctx, "failed to get remote client for shared pool", zap.String("dataplaneName", sharedPool.GetDataplaneName()), zap.Error(err))
		return fmt.Errorf("failed to get remote client for shared pool: %w", err)
	}

	for _, comp := range sharedPool.GetComponentGroups() {

		// get update options
		optionsStr, _ := args.GetString(entity.GetUpdateComponentKey(comp))

		if optionsStr == "" {
			tilog.Info(ctx, "no update options for component, skip", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
			continue
		}

		// Note: comp.Phase comes from the serialized args snapshot (line 100), not from a fresh repo load.
		// On retry, phases remain as the original MODIFYING/CREATING values, so this AVAILABLE guard is safe
		// and will not incorrectly skip components that became AVAILABLE during the first attempt.
		if comp.GetPhase() == apiv1.Component_AVAILABLE {
			tilog.Info(ctx, "component is already available, skip update", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
			continue
		}

		var options []tioperator.Option
		for _, optStr := range strings.Split(optionsStr, ",") {
			options = append(options, tioperator.Option(optStr))
		}

		tilog.Info(ctx, "update component", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Any("options", options))

		err := a.updateComp(ctx, args, remoteCli, sharedPool, comp, options)
		if err != nil {
			tilog.Error(ctx, "failed to update component", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
			return fmt.Errorf("failed to update component: %w", err)
		}

		tilog.Info(ctx, "update component success", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
	}

	return nil
}

func (a *UpdateComponentGroupAction) updateComp(
	ctx context.Context,
	args actions.ActionArgs,
	remoteCli client.Client,
	sharedPool *entity.SharedPool,
	comp *entity.Component,
	options []tioperator.Option,
) error {

	var updateOptions []tioperator.UpdateOption

	for _, opt := range options {
		switch opt {
		case tioperator.OpsUpdateNodeCount:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetNodeQuota().GetNodeCount()})
		case tioperator.OpsUpdateNodeSpec:
			nodeSpec := comp.GetNodeQuota().GetNodeSpec()
			if nodeSpec == nil {
				return fmt.Errorf("node spec is nil for component %s (type: %s)", comp.GetId(), entity.GetComponentTypeName(comp))
			}
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: nodeSpec})
		case tioperator.OpsUpdateMaxCount:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetNodeQuota().GetMaxNodeCount()})
		case tioperator.OpsUpdateMinCount:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetNodeQuota().GetMinNodeCount()})
		case tioperator.OpsUpdateVersion:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetVersion()})
		case tioperator.OpsUpdateImage:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: fmt.Sprintf("%s/%s/%s:%s",
				a.imageConfig.Registry,
				a.imageConfig.RepoPrefix,
				path.Base(comp.GetConfiguration().GetImageRepo()),
				comp.GetConfiguration().GetImageTag(),
			)})
		case tioperator.OpsUpdateConfig:
			if comp.GetComponentType() == apiv1.ComponentType_TIFLASH_LEARNER || comp.GetComponentType() == apiv1.ComponentType_TIFLASH_COMPUTE_LEARNER {
				// for tiflash learner, need to load tiflash config from cluster config
				tiflashLeaderConfig, err := getTiFlashLearnerConfig(ctx, a.clusterConfigClient, comp, sharedPool)
				if err != nil {
					tilog.Error(ctx, "failed to get tiflash leader config", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
					return fmt.Errorf("failed to get %s config: %w", entity.GetComponentTypeName(comp), err)
				}
				updateOptions = append(updateOptions, tioperator.UpdateOption{Option: tioperator.OpsUpdateTiFlashLeaderConfig, Value: tiflashLeaderConfig})
			}

			// tiproxy need set proxy.public_enpoints
			if comp.GetComponentType() == apiv1.ComponentType_TIPROXY {
				err := setTiProxyPublicEndpointConfig(ctx, sharedPool, comp, a.tidbGroupRepo, a.lbClient, a.dataplaneClient)
				if err != nil {
					tilog.Error(ctx, "failed to set tiproxy public endpoint config", zap.String("componentId", comp.GetId()), zap.Error(err))
					return fmt.Errorf("failed to set tiproxy public endpoint config: %w", err)
				}
			}

			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetConfiguration().GetConfigurationFileContent()})
		case tioperator.OpsUpdateUpdateStrategy:
			strategyStr, err := args.GetString(workflow.UpdateStrategyKey)
			if err != nil {
				tilog.Error(ctx, "failed to get update strategy from args", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
				return fmt.Errorf("failed to get update strategy from args: %w", err)
			}
			if strategyStr == "" {
				tilog.Error(ctx, "update strategy is empty", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
				return fmt.Errorf("update strategy is empty for component %s", comp.GetId())
			}
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: v1alpha1.ConfigUpdateStrategy(strategyStr)})

		case tioperator.OpsUpdateTiDBAuditLog:
			// load tidb group
			tidbGroup, err := a.tidbGroupRepo.Load(ctx, comp.GetId(), false)
			if err != nil {
				return fmt.Errorf("failed to load tidb group: %w", err)
			}
			// get tidb cluster
			tidbCluster, err := a.tidbClusterRepo.Load(ctx, tidbGroup.GetClusterId(), false)
			if err != nil {
				return fmt.Errorf("failed to load tidb cluster: %w", err)
			}

			auditLogConfigYaml, err := genAuditlogConfig(ctx, comp, tidbCluster, sharedPool)
			if err != nil {
				tilog.Error(ctx, "failed to generate audit log config", zap.String("componentId", comp.GetId()), zap.Error(err))
				return fmt.Errorf("failed to generate audit log config: %w", err)
			}

			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: auditLogConfigYaml})
		case tioperator.OpsUpdateLabelCidr:
			// load tidb group
			tidbGroup, err := a.tidbGroupRepo.Load(ctx, comp.GetId(), false)
			if err != nil {
				return fmt.Errorf("failed to load tidb group: %w", err)
			}

			lbcidrs, err := getTiDBCidrs(ctx, a.lbClient, tidbGroup, "")
			if err != nil {
				tilog.Error(ctx, "failed to get load balancer cidrs", zap.Error(err))
				return fmt.Errorf("failed to get load balancer cidrs: %w", err)
			}

			if len(lbcidrs) > 0 {
				updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: strings.Join(lbcidrs, ",")})
			}
		case tioperator.OpsUpdateLabels:
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: comp.GetLabels()})
		case tioperator.OpsUpdateTiFlashLeaderConfig:
			tiflashLeaderConfig, err := getTiFlashLearnerConfig(ctx, a.clusterConfigClient, comp, sharedPool)
			if err != nil {
				tilog.Error(ctx, "failed to get tiflash leader config", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
				return fmt.Errorf("failed to get %s config: %w", entity.GetComponentTypeName(comp), err)
			}
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: tiflashLeaderConfig})
		case tioperator.OpsUpdateDataSize:
			for _, storage := range comp.GetStorages() {
				if storage.GetName() == apiv1.Component_Storage_DATA {
					updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: storage.GetSizeGi()})
				}
			}
		case tioperator.OpsUpdateDataVolumeAttributesClass:
			for _, storage := range comp.GetStorages() {
				if storage.GetName() == apiv1.Component_Storage_DATA {
					updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: workflow.GenVolumeAttributesClassName(storage)})
				}
			}
		case tioperator.OpsUpdateRaftLogSize:
			for _, storage := range comp.GetStorages() {
				if storage.GetName() == apiv1.Component_Storage_RAFTLOG {
					updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: storage.GetSizeGi()})
				}
			}
		case tioperator.OpsUpdateRaftLogVolumeAttributesClass:
			for _, storage := range comp.GetStorages() {
				if storage.GetName() == apiv1.Component_Storage_RAFTLOG {
					updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: workflow.GenVolumeAttributesClassName(storage)})
				}
			}
		case tioperator.OpsUpdateNodeVersion:
			nodeVersion, err := args.GetString(workflow.NodeVersionKey)
			if err != nil {
				tilog.Error(ctx, "failed to get node version from args", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
				return fmt.Errorf("failed to get node version from args: %w", err)
			}
			if nodeVersion == "" {
				tilog.Error(ctx, "node version is empty", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
				return fmt.Errorf("node version is empty for component %s", comp.GetId())
			}
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: nodeVersion})
		case tioperator.OpsUpdateTiKVRemoteWorkers:
			target, err := getTiKVRemoteWorkerTargetFromArgs(ctx, args)
			if err != nil {
				tilog.Error(ctx, "failed to get TiKV remote worker target from args", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)), zap.Error(err))
				return fmt.Errorf("failed to get TiKV remote worker target from args: %w", err)
			}
			updateOptions = append(updateOptions, tioperator.UpdateOption{Option: opt, Value: target})
		}
	}

	if len(updateOptions) == 0 {
		tilog.Info(ctx, "no valid update options for component, skip", zap.String("componentId", comp.GetId()), zap.String("componentType", entity.GetComponentTypeName(comp)))
		return nil
	}

	var err error
	switch comp.GetComponentType() {
	case apiv1.ComponentType_PD:
		_, err = tioperator.UpdatePDGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_PD_TSO:
		_, err = tioperator.UpdateTsoGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_PD_SCHEDULING:
		_, err = tioperator.UpdateSchedulingGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_RESOURCE_MANAGER:
		_, err = tioperator.UpdateResourceManagerGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIDB:
		_, err = tioperator.UpdateTiDBGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIDB_WORKER:
		_, err = tioperator.UpdateTiDBWorker(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIKV:
		_, err = tioperator.UpdateTiKVGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIKV_WORKER:
		_, err = tioperator.UpdateTiKVWorker(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_COPROCESSOR_WORKER:
		_, err = tioperator.UpdateCoprocessorWorker(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_SCHEMA_MANAGER:
		_, err = tioperator.UpdateSchemaManager(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIFLASH_COMPUTE, apiv1.ComponentType_TIFLASH_COMPUTE_LEARNER:
		_, err = tioperator.UpdateTiFlashCN(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIFLASH_WRITE, apiv1.ComponentType_TIFLASH_LEARNER:
		_, err = tioperator.UpdateTiFlashWN(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	case apiv1.ComponentType_TIPROXY:
		_, err = tioperator.UpdateTiProxyGroup(ctx, remoteCli, sharedPool.GetId(), comp.GetId(), updateOptions)
	default:
		return fmt.Errorf("unsupported component type: %s", comp.GetComponentType().String())
	}

	if err != nil {
		return fmt.Errorf("failed to update component: %w", err)
	}
	return nil
}

func getTiKVRemoteWorkerTargetFromArgs(ctx context.Context, args actions.ActionArgs) (tioperator.TiKVRemoteWorkerTarget, error) {
	// Get shared pool snapshot from args
	sharedPoolStr, err := args.GetString(workflow.SharedPoolKey)
	if err != nil {
		return tioperator.TiKVRemoteWorkerTarget{}, fmt.Errorf("failed to get sharedPool from args: %w", err)
	}

	var sharedPool entity.SharedPool
	if err := protojson.Unmarshal([]byte(sharedPoolStr), &sharedPool); err != nil {
		return tioperator.TiKVRemoteWorkerTarget{}, fmt.Errorf("failed to unmarshal sharedPool: %w", err)
	}

	// Use shared resolver to compute target from DAG snapshot
	resolved, err := entity.ResolveTiKVRemoteWorkerGroups(&sharedPool)
	if err != nil {
		tilog.Error(ctx, "failed to resolve TiKV remote worker groups", zap.Error(err))
		return tioperator.TiKVRemoteWorkerTarget{}, fmt.Errorf("failed to resolve TiKV remote worker groups: %w", err)
	}

	// Convert to operator target
	return tioperator.TiKVRemoteWorkerTarget{
		Worker:      resolved.Worker,
		Coprocessor: resolved.Coprocessor,
	}, nil
}