package server

import (
	"context"
	"fmt"

	"github.com/samber/lo"
	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/storage/repo"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.opentelemetry.io/otel"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

var dmClusterTracer = otel.Tracer("DMClusterService")

func (s *serverImpl) CreateDMCluster(ctx context.Context, dmCluster *entity.DMCluster) (*entity.DMCluster, error) {
	if dmCluster.GetId() == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dm cluster id is required")
	}

	// Validate dm_cluster_id format
	if err := entity.ValidateLowercaseAlphanumericId("dm_cluster_id", dmCluster.GetId()); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "%s", err.Error())
	}
	if _, err := entity.DMGroupNameV2(dmCluster.GetId()); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid dm cluster id for operator v2: %s", err)
	}

	if dmCluster.GetSharedPoolId() == "" {
		return nil, status.Errorf(codes.InvalidArgument, "shared pool id is required")
	}

	if dmCluster.GetMasterNodeQuota() == nil {
		return nil, status.Errorf(codes.InvalidArgument, "master node quota is required")
	}

	if dmCluster.GetWorkerNodeQuota() == nil {
		return nil, status.Errorf(codes.InvalidArgument, "worker node quota is required")
	}

	var dagExecId string
	var dagSubmitted bool
	// DM clusters are managed by TiDB Operator v2. Keep the backend field on
	// the API/model for compatibility with existing records, but do not expose
	// a legacy deployment path for new clusters.
	backend := entity.DMBackendOperatorV2Proto
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		// check if dm cluster already exists
		existingCluster, err := s.dmClusterRepo.Load(ctx, dmCluster.GetId(), true)
		if err == nil && existingCluster != nil {
			return status.Errorf(codes.AlreadyExists, "dm cluster %s already exists", dmCluster.GetId())
		}

		// load shared pool
		sharedPool, err := s.sharedPoolRepo.Load(ctx, dmCluster.GetSharedPoolId(), false)
		if err != nil {
			tilog.Error(ctx, "failed to load shared pool", zap.String("sharedPoolId", dmCluster.GetSharedPoolId()), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to load shared pool: %s", err)
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return err
		}

		// check shared pool status
		if sharedPool.GetPhase() == apiv1.SharedPool_UNAVAILABLE {
			return status.Errorf(codes.FailedPrecondition, "shared pool is unavailable")
		}

		// set dm cluster dataplane name
		dmCluster.DataplaneName = sharedPool.GetDataplaneName()
		dmCluster.Phase = apiv1.DMCluster_CREATING
		dmCluster.Backend = backend

		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)
		}

		// validate and fill master node spec
		if err := s.validateAndFillDMNodeSpec(ctx, dmCluster.MasterNodeQuota, nodeSpecMap); err != nil {
			tilog.Error(ctx, "failed to validate master node spec", zap.Error(err))
			return err
		}

		// validate and fill worker node spec
		if err := s.validateAndFillDMNodeSpec(ctx, dmCluster.WorkerNodeQuota, nodeSpecMap); err != nil {
			tilog.Error(ctx, "failed to validate worker node spec", zap.Error(err))
			return err
		}

		// save dm cluster
		if err := s.dmClusterRepo.Save(ctx, dmCluster); err != nil {
			tilog.Error(ctx, "failed to save dm cluster", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to save dm cluster: %s", err)
		}

		// build dag args
		dmClusterStr, err := protojson.Marshal(dmCluster)
		if err != nil {
			tilog.Error(ctx, "failed to marshal dm cluster", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to marshal dm cluster: %s", err)
		}

		dagArgs := actions.ActionArgs{
			workflow.SharedPoolIdKey:  actions.ActionArgString(dmCluster.GetSharedPoolId()),
			workflow.DMClusterIdKey:   actions.ActionArgString(dmCluster.GetId()),
			workflow.DMClusterKey:     actions.ActionArgString(string(dmClusterStr)),
			workflow.DMBackendKey:     actions.ActionArgString(entity.DMBackendOperatorV2),
			workflow.DMUpdateImageKey: actions.ActionArgBool(true),
			utils.TraceIdKey:          utils.GetTraceId(ctx),
		}

		// submit dag
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("create-dmcluster-%s", dmCluster.GetId()))
		err = s.dagExecutor.Submit(ctx, dag.PlanApplyDMCluster, dagArgs, dagExecId)
		if err != nil {
			tilog.Error(ctx, "failed to submit dag", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to submit dag: %s", err)
		}
		dagSubmitted = true
		tilog.Info(ctx, "submit create dm cluster dag success", 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", zap.String("dagExecId", dagExecId), zap.Error(abortErr))
			}
		}
		return nil, err
	}

	return dmCluster, nil
}

func (s *serverImpl) GetDMCluster(ctx context.Context, sharedPoolId, dmClusterId string) (*entity.DMCluster, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "shared pool id is required")
	}

	if dmClusterId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dm cluster id is required")
	}

	dmCluster, err := s.dmClusterRepo.Load(ctx, dmClusterId, false)
	if err != nil {
		if status.Code(err) != codes.NotFound {
			return nil, status.Errorf(codes.Internal, "failed to load dm cluster: %s", err)
		}
		return nil, status.Errorf(codes.NotFound, "dm cluster %s not found", dmClusterId)
	}

	if dmCluster.GetSharedPoolId() != sharedPoolId {
		return nil, status.Errorf(codes.PermissionDenied, "dm cluster %s does not belong to shared pool %s", dmClusterId, sharedPoolId)
	}

	// merge status from running dag
	dagIns, err := s.getRunningDagIns(ctx, []string{dag.PlanApplyDMCluster}, []string{sharedPoolId}, nil, nil, nil)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get running dag instances: %s", err)
	}

	for _, ins := range dagIns {
		dmClusterIdFromDag, err := ins.ExecTask.Args.GetString(workflow.DMClusterIdKey)
		if err != nil || dmClusterIdFromDag != dmClusterId {
			continue
		}

		dmClusterStr, err := ins.ExecTask.Args.GetString(workflow.DMClusterKey)
		if err != nil {
			tilog.Warn(ctx, "failed to get dm cluster from exec task args", zap.String("task name", ins.ExecTask.ExecId))
			continue
		}

		var tempDMCluster entity.DMCluster
		err = protojson.Unmarshal([]byte(dmClusterStr), &tempDMCluster)
		if err != nil {
			tilog.Error(ctx, "failed to unmarshal dm cluster", zap.Error(err))
			continue
		}

		dmCluster.MasterNodeQuota = tempDMCluster.GetMasterNodeQuota()
		dmCluster.WorkerNodeQuota = tempDMCluster.GetWorkerNodeQuota()
		dmCluster.Phase = tempDMCluster.GetPhase()
	}

	return dmCluster, nil
}

func (s *serverImpl) ListDMClusters(ctx context.Context, sharedPoolId string) ([]*entity.DMCluster, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "shared pool id is required")
	}

	dmClusters, err := s.dmClusterRepo.List(ctx, &repo.DMClusterListFilter{
		SharedPoolIds: []string{sharedPoolId},
		Deleted:       false,
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list dm clusters: %s", err)
	}

	// merge status from running dags
	dagIns, err := s.getRunningDagIns(ctx, []string{dag.PlanApplyDMCluster}, []string{sharedPoolId}, nil, nil, nil)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get running dag instances: %s", err)
	}

	dmClustersMap := make(map[string]*entity.DMCluster)
	for _, dmCluster := range dmClusters {
		dmClustersMap[dmCluster.GetId()] = dmCluster
	}

	for _, ins := range dagIns {
		dmClusterIdFromDag, err := ins.ExecTask.Args.GetString(workflow.DMClusterIdKey)
		if err != nil {
			tilog.Warn(ctx, "failed to get dm cluster id from exec task args", zap.String("task name", ins.ExecTask.ExecId))
			continue
		}

		dmClusterStr, err := ins.ExecTask.Args.GetString(workflow.DMClusterKey)
		if err != nil {
			tilog.Warn(ctx, "failed to get dm cluster from exec task args", zap.String("task name", ins.ExecTask.ExecId))
			continue
		}
		var tempDMCluster entity.DMCluster
		err = protojson.Unmarshal([]byte(dmClusterStr), &tempDMCluster)
		if err != nil {
			tilog.Error(ctx, "failed to unmarshal dm cluster", zap.String("dmClusterId", dmClusterIdFromDag), zap.String("task name", ins.ExecTask.ExecId), zap.Error(err))
			continue
		}
		if dmCluster, ok := dmClustersMap[dmClusterIdFromDag]; ok {
			dmCluster.MasterNodeQuota = tempDMCluster.GetMasterNodeQuota()
			dmCluster.WorkerNodeQuota = tempDMCluster.GetWorkerNodeQuota()
			dmCluster.Phase = tempDMCluster.GetPhase()
		} else {
			tilog.Warn(ctx, "dm cluster from running dag not found in db", zap.String("dmClusterId", dmClusterIdFromDag))
			dmClustersMap[dmClusterIdFromDag] = &tempDMCluster
		}
	}

	return lo.Values(dmClustersMap), nil
}

func (s *serverImpl) ScaleDMCluster(ctx context.Context, sharedPoolId, dmClusterId string, masterNodeQuota, workerNodeQuota *entity.NodeQuota, force bool) (*entity.DMCluster, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "shared pool id is required")
	}

	if dmClusterId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dm cluster id is required")
	}

	// Ensure DAG exec id generation has a real trace id; background callers may pass a context without one.
	ctx, span := dmClusterTracer.Start(ctx, "ScaleDMCluster")
	defer span.End()

	var dmCluster *entity.DMCluster
	var dagExecId string
	var dagSubmitted bool

	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		var err error
		dmCluster, err = s.dmClusterRepo.Load(ctx, dmClusterId, false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load dm cluster: %s", err)
			}
			return status.Errorf(codes.NotFound, "dm cluster %s not found", dmClusterId)
		}

		if dmCluster.GetSharedPoolId() != sharedPoolId {
			return status.Errorf(codes.PermissionDenied, "dm cluster %s does not belong to shared pool %s", dmClusterId, sharedPoolId)
		}

		// check dm cluster phase
		if !force && dmCluster.GetPhase() != apiv1.DMCluster_AVAILABLE {
			return status.Errorf(codes.FailedPrecondition, "dm cluster is not available, current phase: %s", dmCluster.GetPhase().String())
		}

		// merge status from running dag
		dagIns, err := s.getRunningDagIns(ctx, []string{dag.PlanApplyDMCluster}, []string{sharedPoolId}, nil, nil, nil)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get running dag instances: %s", err)
		}

		if len(dagIns) > 0 && !force {
			return status.Errorf(codes.FailedPrecondition, "there is a running operation on dm cluster, please try again later or set force to true")
		}

		// load shared pool to get node spec map
		sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
		if err != nil {
			tilog.Error(ctx, "failed to load shared pool", zap.String("sharedPoolId", sharedPoolId), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to load shared pool: %s", err)
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return err
		}

		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)
		}

		// update node quota
		targetDMCluster := &entity.DMCluster{
			Id:              dmClusterId,
			SharedPoolId:    sharedPoolId,
			DataplaneName:   dmCluster.GetDataplaneName(),
			MasterNodeQuota: dmCluster.GetMasterNodeQuota(),
			WorkerNodeQuota: dmCluster.GetWorkerNodeQuota(),
		}

		// check worker node first
		if workerNodeQuota != nil {
			targetDMCluster.GetWorkerNodeQuota().NodeCount = workerNodeQuota.GetNodeCount()
			if workerNodeQuota.GetNodeSpec() != nil {
				targetDMCluster.GetWorkerNodeQuota().NodeSpec = workerNodeQuota.GetNodeSpec()
			}

			if err := s.validateAndFillDMNodeSpec(ctx, targetDMCluster.GetWorkerNodeQuota(), nodeSpecMap); err != nil {
				tilog.Error(ctx, "failed to validate worker node spec", zap.Error(err))
				return err
			}
		}

		if masterNodeQuota != nil {
			if masterNodeQuota.GetNodeCount() == 0 && dmCluster.GetWorkerNodeQuota().GetNodeCount() != 0 {
				// master node count can only be scaled to 0 when the current worker node count is 0
				return status.Errorf(codes.InvalidArgument, "only when the current worker node count is 0, the master node count can be scaled to 0")
			}
			targetDMCluster.GetMasterNodeQuota().NodeCount = masterNodeQuota.GetNodeCount()
			if masterNodeQuota.GetNodeSpec() != nil {
				targetDMCluster.GetMasterNodeQuota().NodeSpec = masterNodeQuota.GetNodeSpec()
			}

			// validate and fill node specs
			if err := s.validateAndFillDMNodeSpec(ctx, targetDMCluster.GetMasterNodeQuota(), nodeSpecMap); err != nil {
				tilog.Error(ctx, "failed to validate master node spec", zap.Error(err))
				return err
			}
		}

		// update phase
		dmCluster.Phase = apiv1.DMCluster_MODIFYING
		if err := s.dmClusterRepo.Save(ctx, dmCluster); err != nil {
			tilog.Error(ctx, "failed to save dm cluster", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to save dm cluster: %s", err)
		}

		backend, err := s.dmClusterRepo.GetBackend(ctx, dmClusterId, false)
		if err != nil {
			tilog.Error(ctx, "failed to load dm cluster backend", zap.String("dmClusterId", dmClusterId), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to load dm cluster backend: %s", err)
		}

		// build dag args
		dmClusterStr, err := protojson.Marshal(targetDMCluster)
		if err != nil {
			tilog.Error(ctx, "failed to marshal dm cluster", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to marshal dm cluster: %s", err)
		}

		dagArgs := actions.ActionArgs{
			workflow.SharedPoolIdKey: actions.ActionArgString(sharedPoolId),
			workflow.DMClusterIdKey:  actions.ActionArgString(dmClusterId),
			workflow.DMClusterKey:    actions.ActionArgString(string(dmClusterStr)),
			workflow.DMBackendKey:    actions.ActionArgString(backend),
			utils.TraceIdKey:         utils.GetTraceId(ctx),
		}

		// submit dag
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("scale-dmcluster-%s", dmClusterId))
		err = s.dagExecutor.Submit(ctx, dag.PlanApplyDMCluster, dagArgs, dagExecId)
		if err != nil {
			tilog.Error(ctx, "failed to submit dag", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to submit dag: %s", err)
		}
		dagSubmitted = true
		tilog.Info(ctx, "submit scale dm cluster dag success", 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", zap.String("dagExecId", dagExecId), zap.Error(abortErr))
			}
		}
		return nil, err
	}

	// return target dm cluster with updated quota
	if masterNodeQuota != nil {
		dmCluster.MasterNodeQuota = masterNodeQuota
	}
	if workerNodeQuota != nil {
		dmCluster.WorkerNodeQuota = workerNodeQuota
	}

	return dmCluster, nil
}

func (s *serverImpl) DeleteDMCluster(ctx context.Context, sharedPoolId, dmClusterId string) error {
	if sharedPoolId == "" {
		return status.Errorf(codes.InvalidArgument, "shared pool id is required")
	}

	if dmClusterId == "" {
		return status.Errorf(codes.InvalidArgument, "dm cluster id is required")
	}

	var dagExecId string
	var dagSubmitted bool

	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		dmCluster, err := s.dmClusterRepo.Load(ctx, dmClusterId, false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load dm cluster: %s", err)
			}
			return status.Errorf(codes.NotFound, "dm cluster %s not found", dmClusterId)
		}

		if dmCluster.GetSharedPoolId() != sharedPoolId {
			return status.Errorf(codes.PermissionDenied, "dm cluster %s does not belong to shared pool %s", dmClusterId, sharedPoolId)
		}
		if err := s.guardSharedPoolMutationByID(ctx, sharedPoolId, sharedPoolMutationDeleteInstance); err != nil {
			return err
		}

		// cancel any running dag
		runningDagIns, err := s.getRunningDagIns(ctx, []string{dag.PlanApplyDMCluster}, []string{sharedPoolId}, nil, nil, nil)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get running dag instances: %s", err)
		}

		for id, ins := range runningDagIns {
			dmClusterIdFromDag, err := ins.ExecTask.Args.GetString(workflow.DMClusterIdKey)
			if err != nil || dmClusterIdFromDag != dmClusterId {
				continue
			}

			tilog.Info(ctx, "aborting dm cluster dag", zap.String("exec id", id))
			if abortErr := s.dagExecutor.Abort(ctx, ins.ExecTask.ExecId); abortErr != nil {
				tilog.Error(ctx, "failed to abort dm cluster dag", zap.String("task id", ins.ExecTask.ExecId), zap.Error(abortErr))
				return status.Errorf(codes.Internal, "failed to abort dm cluster dag: %s", abortErr)
			}
		}

		if err := s.dmClusterRepo.Delete(ctx, dmClusterId); err != nil {
			tilog.Error(ctx, "failed to delete dm cluster from db", zap.String("dmClusterId", dmClusterId), zap.Error(err))
			return status.Errorf(codes.Internal, "failed to delete dm cluster: %s", err)
		}

		// build dag args
		dmClusterStr, err := protojson.Marshal(dmCluster)
		if err != nil {
			tilog.Error(ctx, "failed to marshal dm cluster", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to marshal dm cluster: %s", err)
		}

		dagArgs := actions.ActionArgs{
			workflow.SharedPoolIdKey: actions.ActionArgString(sharedPoolId),
			workflow.DMClusterIdKey:  actions.ActionArgString(dmClusterId),
			workflow.DMClusterKey:    actions.ActionArgString(string(dmClusterStr)),
			workflow.DMBackendKey:    actions.ActionArgString(entity.DMBackendOperatorV2),
			workflow.DMDeleteKey:     actions.ActionArgString(workflow.DMDeleteValue),
			utils.TraceIdKey:         utils.GetTraceId(ctx),
		}

		// submit delete dag
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("delete-dmcluster-%s", dmClusterId))
		err = s.dagExecutor.Submit(ctx, dag.PlanApplyDMCluster, dagArgs, dagExecId)
		if err != nil {
			tilog.Error(ctx, "failed to submit delete dag", zap.Error(err))
			return status.Errorf(codes.Internal, "failed to submit delete dag: %s", err)
		}
		dagSubmitted = true
		tilog.Info(ctx, "submit delete dm cluster dag success", 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", zap.String("dagExecId", dagExecId), zap.Error(abortErr))
			}
		}
		return err
	}

	return nil
}

// validateAndFillDMNodeSpec validates and fills the node spec for DM cluster
func (s *serverImpl) validateAndFillDMNodeSpec(ctx context.Context, nodeQuota *entity.NodeQuota, nodeSpecMap map[string][]*entity.NodeSpec) error {
	if nodeQuota == nil {
		return status.Errorf(codes.InvalidArgument, "node quota is required")
	}

	if nodeQuota.GetNodeSpec() == nil {
		return status.Errorf(codes.InvalidArgument, "node spec is required")
	}

	// DM components don't use local disk
	fullNodeSpec, err := getNodeSpec(ctx, nodeQuota.GetNodeSpec(), nodeSpecMap, false)
	if err != nil {
		return status.Errorf(codes.InvalidArgument, "invalid node spec: %s", err)
	}

	// fill the full node spec
	nodeQuota.NodeSpec = fullNodeSpec

	return nil
}