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

import (
	"context"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/tiflow/dm/openapi"
	"github.com/tidbcloud/dataflow-service/internal/db"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/internal/handler/agent"
	"github.com/tidbcloud/dataflow-service/internal/metrics"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"github.com/tidbcloud/eks-provider/api/v1alpha1"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

type StateDump struct {
	StateBase
}

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

func (s *StateDump) 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

	if !m.needDumpStage(dmJob) {
		logger.Info("skip dump phase because this dm job not required it")
		next = new(StateLoad)
		return
	}

	jobSources, err := GetDMJobSources(dmJob.ID)
	if err != nil {
		logger.Error("failed get dm job source", zap.Error(err))
		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
	}
	for _, jobSource := range jobSources {
		if jobSource.DumpStatus.IsFinished() {
			logger.Info("skip dump phase, dump task has been completed")
			next = new(StateLoad)
			return
		}
	}

	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))
		return
	} else if dmJobCR.Status.Phase != v1alpha1.DMJobPhaseRunning {
		// this should not happend
		logger.Error("dm job cr not running, this should not happend!", zap.Any("phase", dmJobCR.Status.Phase))
		return
	}

	taskName := generateDMTaskName(dmJob.ID, entity.DMStageDump)
	logger = logger.With(zap.String("taskName", taskName))
	ctx = log.WithLogger(ctx, logger)
	sourceNames := GetDatasourceNames(jobSources)

	// query dump task status, and create dump task if not found
	// masterServerName := "http://fangdong-test-002-dm-master.dm-test.svc.cluster.local:8261"
	masterServerName := dmJobCR.Status.MasterServerName
	dmOpenapiClient := c.agentClient.GetDMOpenapiClient(dmJob.ClusterID.Uint64(), masterServerName)
	taskStatusResp, err := dmOpenapiClient.QueryTaskStatus(ctx, taskName, sourceNames)
	if err == agent.ErrorDMResourceNotFound {
		// create dump task
		_, err = s.createDMTask(c, dmJobCR, openapi.TaskTaskModeDump, taskName, dmOpenapiClient, jobSources)
		if err != nil {
			logger.Error("failed to create dump task", zap.Error(err))
			metrics.IncDMCreateStageFailure(c.job.DFJob, "dump_create", "create_task")
			err = errors.Annotate(err, "failed to create dump task")
			return
		}
		// start task
		if err = dmOpenapiClient.StartTask(ctx, taskName, &openapi.StartTaskRequest{
			SourceNameList: &sourceNames,
		}); err != nil {
			logger.Error("failed to start dump task", zap.Error(err))
			metrics.IncDMCreateStageFailure(c.job.DFJob, "dump_start", "start_task")
			err = errors.Annotate(err, "failed to start dump task")
			return
		}
		return
	} else if err != nil {
		err = errors.Annotate(err, "failed to query dump task status")
		return
	}

	taskStatus := taskStatusResp.Data
	if len(taskStatus) == 0 {
		// should not happend
		err = errors.New("dump task status is empty, this should not happend!")
		return
	}
	for i := range taskStatus {
		if taskStatus[i].Unit != DMTaskUintDump {
			logger.Error("dump task status got unexpected unit", zap.Any("taskStatus", taskStatus), zap.Int("idx", i))
			err = errors.Errorf("dump task status got unexpected unit:%s", taskStatus[i].Unit)
			return
		}
	}

	// save dump status
	if err = s.saveDumpStatus(ctx, taskStatus, jobSources); err != nil {
		err = errors.Annotate(err, "failed to save dump status")
		return
	}

	allFinished, err := s.checkSubTaskStatus(c, dmOpenapiClient, entity.DMStageDump, taskStatus)
	if err == nil && allFinished {
		// delete dump task to prevent it from running again
		// https://pingcap.feishu.cn/wiki/WTDUwRrHYiJAAikTzzYcTMA0nJd
		if err = dmOpenapiClient.DeleteTask(ctx, taskName); err != nil {
			logger.Error("failed to delete dump task", zap.Error(err))
			err = errors.Annotate(err, "failed to delete dump task")
			return
		}
		next = new(StateLoad)
		PublishEvent(ctx, &PublishEventReq{Job: c.job})
	}
	return
}

func (s *StateDump) saveDumpStatus(ctx context.Context, taskStatus []openapi.SubTaskStatus, jobSources []entity.DMJobSource) error {
	logger := log.GetLogger(ctx)

	// Map task statuses by source name for quick lookup
	statusMap := make(map[string]*openapi.SubTaskStatus, len(taskStatus))
	for i := range taskStatus {
		statusMap[taskStatus[i].SourceName] = &taskStatus[i]
	}

	// Perform database transaction to update job sources
	err := db.Transaction(func(tx *gorm.DB) error {
		for _, jobSource := range jobSources {
			dataSourceName := generateDMDataSourceName(&jobSource)
			status := statusMap[dataSourceName]
			if status == nil {
				logger.Error("[saveDumpStatus] No matching status found for job source", zap.String("dataSourceName", dataSourceName))
				continue
			}

			// Initialize DumpStatus if nil
			if jobSource.DumpStatus == nil {
				jobSource.DumpStatus = &entity.DMDumpStatus{BeginTime: time.Now().Unix()}
			}
			jobSource.DumpStatus.UpdateTime = time.Now().Unix()
			jobSource.DumpStatus.ErrorMsg = status.ErrorMsg
			jobSource.DumpStatus.Duration = time.Now().Unix() - jobSource.DumpStatus.BeginTime
			jobSource.DumpStatus.Stage = status.Stage
			if status.Stage == "" {
				jobSource.DumpStatus.Stage = openapi.TaskStageRunning
			}

			// Update DumpStatus based on task status
			if status.DumpStatus != nil && status.ErrorMsg == nil {
				jobSource.DumpStatus.TotalTables = status.DumpStatus.TotalTables
				jobSource.DumpStatus.CompletedTables = status.DumpStatus.CompletedTables
				jobSource.DumpStatus.FinishedBytes = status.DumpStatus.FinishedBytes
				jobSource.DumpStatus.FinishedRows = status.DumpStatus.FinishedRows
				jobSource.DumpStatus.Bps = status.DumpStatus.Bps
				jobSource.DumpStatus.Progress = status.DumpStatus.Progress
			}

			// Update job source in the database
			if err := jobSource.Update(tx, entity.DMJobSourceDBSchema.DumpStatus); err != nil {
				logger.Error("[saveDumpStatus] Failed to update job source dump status", zap.String("dataSourceName", dataSourceName), zap.Any("dumpStatus", status.DumpStatus), zap.Error(err))
				return err
			}
		}
		return nil
	})
	if err != nil {
		logger.Error("[saveDumpStatus] Failed to update DM job source status", zap.Error(err))
		return err
	}

	return nil
}