// Package dm add by blacknc <blacknc@163.com> at 2024/10/31 00:19:09
package dm

import (
	"time"

	"github.com/tidbcloud/dataflow-service/internal/db"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/internal/metrics"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

type StateScale struct{}

func (s *StateScale) String() string {
	return "StateScale"
}

func (s *StateScale) Exec(c *TaskContext) (next TaskState, err error) {
	ctx := c.ctx
	dfJob := c.job.DFJob
	dmJob := c.job.DMJob
	logger := log.GetLogger(ctx)
	m := c.model

	dmJobCR, err := m.resourceAPI.GetDMJob(ctx, dmJob.ProjectID.Uint64(), dmJob.ID.Uint64())
	if err != nil {
		logger.Error("failed to get dm job cr", zap.Error(err))
		metrics.IncDMModifyFailure(dfJob, "load_job", "get_cr")
		return
	}

	jobSources, err := GetDMJobSources(dmJob.ID)
	if err != nil {
		logger.Error("failed get dm job source", zap.Error(err))
		metrics.IncDMModifyFailure(dfJob, "load_job", "get_sources")
		return nil, err
	} else if len(jobSources) == 0 {
		// should not happend
		logger.Error("dm job source is empty, this should not happend!")
		return nil, nil
	}

	mods, err := GetUnfinishedJobModification(dmJob.ID)
	if err != nil {
		logger.Error("fail to get job modifications", zap.Error(err))
		metrics.IncDMModifyFailure(dfJob, "load_spec", "get_modification")
		return
	}
	if len(mods) > 1 {
		// should not happen
		logger.Error("multiple modifying operation")
		return
	}

	mod := mods[0]
	m.addPhaseMetric(dfJob, entity.DFJobPhaseModifying, mod.CreatedAt)

	changeSpec, err := GetJobSpecByID(mod.Param.ToSpecID)
	if err != nil {
		logger.Error("failed to get job spec by id", zap.Any("modID", mod.ID), zap.Any("toSpecID", mod.Param.ToSpecID))
		metrics.IncDMModifyFailure(dfJob, "load_spec", "get_spec")
		return
	}

	dmOpenapiClient := c.agentClient.GetDMOpenapiClient(dmJob.ClusterID.Uint64(), dmJobCR.Status.MasterServerName)
	dataSourceName := generateDMDataSourceName(&jobSources[0])
	toSpec := changeSpec.Name
	done, err := m.scaleWorkerToSpec(ctx, dmOpenapiClient, dmJob, dmJobCR, toSpec, dataSourceName)
	if err != nil {
		logger.Error("failed to scale worker for dm job", zap.Any("modID", mod.ID), zap.String("toSpec", toSpec), zap.String("dataSourceName", dataSourceName), zap.Error(err))
		metrics.IncDMModifyFailure(dfJob, "scale_worker", "scale")
		return nil, err
	} else if !done {
		logger.Info("scaling worker for dm laod task, waiting for the next retry.", zap.Any("modID", mod.ID), zap.String("toSpec", toSpec), zap.String("dataSourceName", dataSourceName))
		return nil, nil
	}

	if err = s.updateJobPhase(dfJob, dmJob, &mod); err != nil {
		logger.Error("fail to update job phase after scaling completed", zap.Error(err))
		metrics.IncDMModifyFailure(dfJob, "update_phase", "update")
		return
	} else {
		logger.Info("succeed to update job phase after scaling completed")
		m.deletePhaseMetric(logger, entity.DFJobPhaseModifying, dfJob)
		PublishEvent(ctx, &PublishEventReq{Job: c.job, Action: DMActionScaleCompleted})
	}

	return
}

func (s *StateScale) updateJobPhase(dfJob *entity.DFJob, dmJob *entity.DMJob, mod *entity.DMJobModifications) error {
	curTs := time.Now()
	mod.FinishedAt = &curTs
	dfJob.Phase = entity.DFJobPhaseCreated
	dmJob.SpecID = mod.Param.ToSpecID
	return db.Transaction(func(tx *gorm.DB) error {
		if err1 := updateDFJobPhaseWithExLock(tx, dfJob, entity.DFJobPhaseModifying, dfJob.Phase); err1 != nil {
			return err1
		}
		if err1 := mod.Update(tx, entity.DMJobModificationsDBSchema.FinishedAt); err1 != nil {
			return err1
		}
		if err1 := dmJob.Update(tx, entity.DMJobDBSchema.SpecID); err1 != nil {
			return err1
		}
		return nil
	})
}