// Package cluster_job provides an implementation of the Model interface
// that handles register and deregister ClusterJobs for tasks in TiDB.
// It uses the gRPC client from the pkg/grpc/grpcli package to communicate with the core service.
package cluster_job

import (
	"context"
	"strconv"

	appconfig "github.com/tidbcloud/dataflow-service-ng/app/config"
	"github.com/tidbcloud/dataflow-service-ng/pkg/grpc/grpcli"
	"github.com/tidbcloud/dataflow-service-ng/pkg/log"
	"github.com/tidbcloud/pkgv2/bizerr"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
)

type Model interface {
	RegisterCDCClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError
	DeregisterCDCClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError
	RegisterImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError
	EnsureImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError
	DeregisterImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError
}

// The clusterJobModelImpl struct provides the implementation of the Model interface.
type clusterJobModelImpl struct {
	mgmtSvcCli grpcli.MgmtSvcGrpcClient
}

func NewClusterJobModelImpl(mgmtSvcCli grpcli.MgmtSvcGrpcClient) Model {
	return &clusterJobModelImpl{
		mgmtSvcCli: mgmtSvcCli,
	}
}

// RegisterCDCClusterJob registers a cluster job for a CDC task.
func (job *clusterJobModelImpl) RegisterCDCClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError {
	return job.register(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_CDC)
}

func (job *clusterJobModelImpl) DeregisterCDCClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError {
	return job.deregister(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_CDC)
}

// RegisterImportClusterJob registers a cluster job for an Import Into task.
func (job *clusterJobModelImpl) RegisterImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError {
	return job.register(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_IMPORT)
}

// EnsureImportClusterJob idempotently registers an Import Into cluster job.
// It is used to backfill jobs created before custom status registration was introduced.
func (job *clusterJobModelImpl) EnsureImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError {
	existed, err := job.isRegistered(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_IMPORT)
	if err != nil {
		return err
	}
	if existed {
		return nil
	}

	registerErr := job.register(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_IMPORT)
	if registerErr == nil {
		return nil
	}

	// The registration may have reached TMS even if the client did not receive
	// the response, or another worker may have won the concurrent registration.
	existed, queryErr := job.isRegistered(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_IMPORT)
	if queryErr == nil && existed {
		return nil
	}
	return registerErr
}

func (job *clusterJobModelImpl) DeregisterImportClusterJob(ctx context.Context, clusterID, taskID uint64) bizerr.BizError {
	return job.deregister(ctx, clusterID, taskID, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_IMPORT)
}

func (job *clusterJobModelImpl) register(
	ctx context.Context,
	clusterID uint64,
	taskID uint64, statusType mgmt.CustomStatusType,
) bizerr.BizError {
	setStatusReq := &mgmt.SetClusterCustomStatusRequest{
		CustomStatus: &mgmt.ClusterCustomStatus{
			ClusterId:  clusterID,
			StatusName: strconv.FormatUint(taskID, appconfig.DecimalBase),
			StatusType: statusType,
			Mutex: &mgmt.ClusterCustomStatus_CustomMutex{
				StatusType: []mgmt.CustomStatusType{},
				ClusterAction: []mgmt.ClusterAction{
					mgmt.ClusterAction_CLUSTER_ACTION_DELETE_CLUSTER,
					mgmt.ClusterAction_CLUSTER_ACTION_PAUSE_CLUSTER,
				},
			},
		},
	}

	_, err := job.mgmtSvcCli.SetClusterCustomStatus(ctx, setStatusReq)
	if err != nil {
		log.GetLogger(ctx).Error("failed to register clusterJob",
			zap.Uint64("clusterID", clusterID),
			zap.Any("SetClusterCustomStatus", setStatusReq),
			zap.Error(err))
		return err.Chain("create cluster job")
	}
	return nil
}

func (job *clusterJobModelImpl) deregister(
	ctx context.Context,
	clusterID uint64,
	taskID uint64,
	statusType mgmt.CustomStatusType,
) bizerr.BizError {
	query := &mgmt.IsClusterCustomStatusExistedRequest{
		ClusterId:  clusterID,
		StatusType: statusType,
	}
	ss, err := job.mgmtSvcCli.IsClusterCustomStatusExisted(ctx, query)
	if err != nil {
		log.GetLogger(ctx).Error("failed to deregister clusterJob",
			zap.Uint64("clusterID", clusterID),
			zap.Any("IsClusterCustomStatusExisted", query),
			zap.Error(err))
		return err.Chain("deregister status")
	}
	if ss.IsExisted {
		for _, status := range ss.ExistedCustomStatuses {
			if status.StatusName != strconv.FormatUint(taskID, appconfig.DecimalBase) {
				continue
			}
			err = job.mgmtSvcCli.DeleteClusterCustomStatus(ctx, &mgmt.DeleteClusterCustomStatusRequest{
				ClusterId: clusterID,
				StatusId:  status.StatusId,
			})
			if err != nil {
				log.GetLogger(ctx).Error("failed to deregister clusterJob",
					zap.Uint64("taskID", taskID),
					zap.Uint64("clusterID", clusterID),
				)
				return err.Chain("delete status")
			}
		}
	}
	return nil
}

func (job *clusterJobModelImpl) isRegistered(
	ctx context.Context,
	clusterID uint64,
	taskID uint64,
	statusType mgmt.CustomStatusType,
) (bool, bizerr.BizError) {
	query := &mgmt.IsClusterCustomStatusExistedRequest{
		ClusterId:  clusterID,
		StatusType: statusType,
	}
	ss, err := job.mgmtSvcCli.IsClusterCustomStatusExisted(ctx, query)
	if err != nil {
		return false, err.Chain("query cluster job")
	}
	for _, status := range ss.GetExistedCustomStatuses() {
		if status.GetStatusName() == strconv.FormatUint(taskID, appconfig.DecimalBase) {
			return true, nil
		}
	}
	return false, nil
}