package impl

import (
	"context"
	"fmt"
	"time"

	"github.com/AlekSi/pointer"
	"github.com/tidbcloud/dataflow-service/app/internal/status_matrix"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	"github.com/tidbcloud/dataflow-service/app/models/replay"
	"github.com/tidbcloud/dataflow-service/app/repo"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	"github.com/tidbcloud/dataflow-service/app/utils"
	"github.com/tidbcloud/dataflow-service/internal/config"
	"github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"github.com/tidbcloud/pkgv2/bizerr"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

var Tracer = otel.Tracer("dataflow-worker.replay")

// ActionHandler define the action handler function type
type ActionHandler func(ctx context.Context, replayEntity *entity.Replay) (completed bool, err bizerr.BizError)

type replayWorkerImpl struct {
	mgmtSvcGrpcClient grpcli.MgmtSvcGrpcClient
	replayRepo        repo.ReplayRepo
	replayOperator    replay.ReplayOperator
	statusMatrix      status_matrix.StatusMatrix
	actionHandlers    map[entity.ReplayCurrentState]ActionHandler
}

func NewReplayWorker(
	mgmtSvcGrpcClient grpcli.MgmtSvcGrpcClient,
	replayRepo repo.ReplayRepo,
	replayOperator replay.ReplayOperator,
) replay.ReplayWorker {
	worker := &replayWorkerImpl{
		mgmtSvcGrpcClient: mgmtSvcGrpcClient,
		replayRepo:        replayRepo,
		replayOperator:    replayOperator,
		statusMatrix:      NewReplayStatusMatrix(),
		actionHandlers:    make(map[entity.ReplayCurrentState]ActionHandler),
	}

	// register action handlers
	worker.registerActionHandlers()

	return worker
}

// registerActionHandlers register all action handlers
func (w *replayWorkerImpl) registerActionHandlers() {
	w.actionHandlers[entity.ReplayCurrentStateDeploying] = w.handleDeploying
	w.actionHandlers[entity.ReplayCurrentStateStarting] = w.handleStarting
	w.actionHandlers[entity.ReplayCurrentStateRunning] = w.handleRunning
	w.actionHandlers[entity.ReplayCurrentStatePausing] = w.handlePausing
	w.actionHandlers[entity.ReplayCurrentStateDeleting] = w.handleDeleting
}

// getActionForDesiredState get the action for the desired state (using the status matrix)
func (w *replayWorkerImpl) getActionForDesiredState(desiredState entity.ReplayDesiredState) entity.ReplayActionType {
	switch desiredState {
	case entity.ReplayDesiredStateDeploy:
		return entity.ReplayActionDeploy
	case entity.ReplayDesiredStateActive:
		return entity.ReplayActionActivate
	case entity.ReplayDesiredStatePause:
		return entity.ReplayActionPause
	case entity.ReplayDesiredStateDelete:
		return entity.ReplayActionDelete
	}
	return ""
}

// SyncAllReplays sync all Replay task status, drive the status machine transition
func (w *replayWorkerImpl) SyncAllReplays(ctx context.Context) bizerr.BizError {
	logger := log.GetLogger(ctx)
	logger.Info("syncing all replays")

	// get all Replay tasks that need to be processed (non-terminal states)
	replays, err := w.replayRepo.List(ctx, func(qs entity.ReplayQuerySet) entity.ReplayQuerySet {
		return qs.CurrentStateNotIn(entity.ReplayCurrentStateDeleted)
	})
	if err != nil {
		logger.Error("list replays failed", zap.Error(err))
		return err
	} else if len(replays) == 0 {
		logger.Info("no replays to sync")
		return nil
	}
	logger.Info("found replays to sync", zap.Int("count", len(replays)))

	wp := utils.NewWorkerPoolWrapper(
		utils.WorkerPoolConcurrencyOp(config.Global().ReplayConfig.WorkerConcurrency))

	for i := range replays {
		replayEntity := &replays[i]
		wp.Submit(ctx, func() {
			// create a new span and generate a new trace ID for this sync operation
			oldTraceID := trace.SpanContextFromContext(ctx).TraceID().String()
			ctx, span := Tracer.Start(context.Background(), fmt.Sprintf("sync_one_replay.%d", replayEntity.ID.Uint64()))
			defer span.End()

			ctx = log.WithTraceID(ctx)
			logger := log.GetLogger(ctx).With(
				zap.String("old_trace_id", oldTraceID),
				zap.Uint64("replay_id", uint64(replayEntity.ID)),
			)
			ctx = log.WithLogger(ctx, logger)
			logger.Info("sync replay")
			if err := w.syncOneReplay(ctx, replayEntity); err != nil {
				logger.Error("sync replay failed", zap.Error(err))
			} else {
				logger.Info("sync replay completed")
			}
		})
	}
	wp.StopWait()
	return nil
}

// checkTargetCluster checks if the target cluster exists, returns the action to perform
func (w *replayWorkerImpl) checkTargetCluster(
	ctx context.Context,
	replayEntity *entity.Replay,
	action entity.ReplayActionType,
) (entity.ReplayActionType, bizerr.BizError) {
	if action == entity.ReplayActionDelete {
		return action, nil
	}

	logger := log.GetLogger(ctx)
	_, bizErr := w.mgmtSvcGrpcClient.GetClusterByID(ctx, replayEntity.TargetClusterID.Uint64())
	if bizErr == nil {
		return action, nil
	}

	if bizErr.ErrCode() != bizerr.ErrCodeRecordNotFound {
		return action, bizErr.Chain("get target cluster info")
	}

	if replayEntity.DeletingAt == nil {
		replayEntity.DeletingAt = pointer.ToTime(time.Now().UTC())
	}
	logger.Info("target cluster is deleted, deleting the replay")
	return entity.ReplayActionDelete, nil
}

// syncOneReplay sync one Replay task status
func (w *replayWorkerImpl) syncOneReplay(ctx context.Context, replayEntity *entity.Replay) bizerr.BizError {
	logger := log.GetLogger(ctx)
	// if the replay is in the terminal state, no need to process
	if replayEntity.IsInTerminalState() {
		logger.Info("replay is in terminal state, no need to process")
		return nil
	}
	// get the action for the desired state and current state
	action := w.getActionForDesiredState(replayEntity.DesiredState)
	if action == "" {
		logger.Error("no action for desired state",
			zap.String("desired_state", string(replayEntity.DesiredState)))
		return bizerr.New(bizerr.ErrCodeParams, "no action for desired state")
	}

	action, bizErr := w.checkTargetCluster(ctx, replayEntity, action)
	if bizErr != nil {
		return bizErr.Chain("check target cluster")
	}

	oldStates := []entity.ReplayCurrentState{replayEntity.CurrentState}
	currentStatus := status_matrix.CurrentStatus{Status: string(replayEntity.CurrentState)}
	nextStatus := w.statusMatrix.CalculateNextStatusForAction(currentStatus, string(action))
	if nextStatus != nil {
		replayEntity.CurrentState = entity.ReplayCurrentState(nextStatus.Status)
	}
	// get the corresponding handler
	handler, ok := w.actionHandlers[replayEntity.CurrentState]
	if !ok {
		logger.Info("no handler for current state, nothing to do",
			zap.String("current_state", string(replayEntity.CurrentState)))
		return nil
	}
	// execute the handler, check if the operation is completed
	logger.Info("executing replay handler",
		zap.String("current_state", string(replayEntity.CurrentState)),
		zap.String("desired_state", string(replayEntity.DesiredState)))
	completed, err := handler(ctx, replayEntity)
	if err != nil {
		return err
	}

	w.handleCompletedAction(ctx, replayEntity, &oldStates, completed)

	if err := w.replayRepo.Save(ctx, replayEntity); err != nil {
		return err.Chain("save replay")
	}
	logger.Info("transition state completed",
		zap.Any("old_states", oldStates),
		zap.String("new_state", string(replayEntity.CurrentState)),
		zap.String("action", string(action)))
	return nil
}

// handleCompletedAction handles the completed action and updates the replay state
func (w *replayWorkerImpl) handleCompletedAction(
	ctx context.Context,
	replayEntity *entity.Replay,
	oldStates *[]entity.ReplayCurrentState,
	completed bool,
) {
	if !completed {
		return
	}

	logger := log.GetLogger(ctx)
	nextStatus := w.statusMatrix.CalculateNextStatusIfActionDone(
		status_matrix.CurrentStatus{Status: string(replayEntity.CurrentState)}, "")
	logger.Info("operation completed, calculating next status",
		zap.String("current_state", string(replayEntity.CurrentState)),
		zap.Any("next_status", nextStatus))

	if nextStatus != nil {
		*oldStates = append(*oldStates, replayEntity.CurrentState)
		replayEntity.CurrentState = entity.ReplayCurrentState(nextStatus.Status)
	}

	if replayEntity.CurrentState == entity.ReplayCurrentStateDeleted && replayEntity.DeletedAt.Time.IsZero() {
		replayEntity.DeletedAt = gorm.DeletedAt{Time: time.Now(), Valid: true}
	}
}

// ==================== Action Handlers ====================

// handleDeploying handle the deploy operation
func (w *replayWorkerImpl) handleDeploying(ctx context.Context, replayEntity *entity.Replay) (bool, bizerr.BizError) {
	logger := log.GetLogger(ctx)

	// get the status, if the replay is not found, deploy the replay
	statusResp, err := w.replayOperator.GetStatus(ctx, &domain.GetReplayStatusReq{Replay: replayEntity})
	if err != nil {
		if err.ErrCode() != bizerr.ErrCodeRecordNotFound {
			logger.Error("get replay status failed", zap.Error(err))
			return false, err
		}

		// execute the deploy operation
		err := w.replayOperator.Deploy(ctx, &domain.DeployReplayReq{Replay: replayEntity})
		if err != nil {
			logger.Error("deploy replay failed", zap.Error(err))
			return false, err.Chain("deploy replay")
		}
		return false, nil
	}

	if statusResp.Status != nil {
		replayEntity.ReplayStatus = statusResp.Status
	}

	completed := w.isDeployComplete(statusResp)
	if completed {
		logger.Info("deploy action completed")
	}
	return completed, nil
}

// handleStaring handle the activate operation
func (w *replayWorkerImpl) handleStarting(ctx context.Context, replayEntity *entity.Replay) (bool, bizerr.BizError) {
	logger := log.GetLogger(ctx)

	// get the status, if the replay is not active, modify the desired state to active
	statusResp, err := w.replayOperator.GetStatus(ctx, &domain.GetReplayStatusReq{Replay: replayEntity})
	if err != nil {
		logger.Error("get replay status failed", zap.Error(err))
		return false, err.Chain("get replay status")
	} else if statusResp.State != entity.ReplayDesiredStateActive {
		// Modify will change replay task configuration and trigger the start operation
		err := w.replayOperator.Modify(ctx,
			&domain.ModifyReplayReq{Replay: replayEntity, State: entity.ReplayDesiredStateActive})
		if err != nil {
			logger.Error("modify replay failed", zap.Error(err))
			return false, err.Chain("modify replay")
		}
		return false, nil // wait for the next sync
	}

	// update the status information
	if statusResp.Status != nil {
		replayEntity.ReplayStatus = statusResp.Status
	}

	// check if the replay is running
	completed := w.isRunning(statusResp)
	if completed {
		logger.Info("starting action completed")
	}
	return completed, nil
}

// handleRunning handle the running operation
func (w *replayWorkerImpl) handleRunning(
	ctx context.Context, replayEntity *entity.Replay,
) (bool, bizerr.BizError) {
	logger := log.GetLogger(ctx)

	statusResp, err := w.replayOperator.GetStatus(ctx, &domain.GetReplayStatusReq{Replay: replayEntity})
	if err != nil {
		logger.Error("get replay status failed", zap.Error(err))
		return false, err.Chain("get replay status")
	}

	// update the status information
	if statusResp.Status != nil {
		replayEntity.ReplayStatus = statusResp.Status
	}

	return true, nil
}

// handlePausing handle the pause operation
func (w *replayWorkerImpl) handlePausing(ctx context.Context, replayEntity *entity.Replay) (bool, bizerr.BizError) {
	logger := log.GetLogger(ctx)

	// get the status, if the replay is not paused, modify the desired state to paused
	statusResp, err := w.replayOperator.GetStatus(ctx, &domain.GetReplayStatusReq{Replay: replayEntity})
	if err != nil {
		logger.Error("get replay status failed", zap.Error(err))
		return false, err.Chain("get replay status")
	} else if statusResp.State != entity.ReplayDesiredStatePause {
		// execute the stop operation
		err := w.replayOperator.Stop(ctx, &domain.StopReplayReq{Replay: replayEntity})
		if err != nil {
			logger.Error("stop replay failed", zap.Error(err))
			return false, err.Chain("stop replay")
		}
		return false, nil // wait for the next sync
	}

	// update the status information
	if statusResp.Status != nil {
		replayEntity.ReplayStatus = statusResp.Status
	}

	// check if the replay is paused
	completed := w.isPaused(statusResp)
	if completed {
		logger.Info("pause action completed")
	}
	return completed, nil
}

// handleDeleting handle the delete operation
func (w *replayWorkerImpl) handleDeleting(ctx context.Context, replayEntity *entity.Replay) (bool, bizerr.BizError) {
	logger := log.GetLogger(ctx)

	statusResp, err := w.replayOperator.GetStatus(ctx, &domain.GetReplayStatusReq{Replay: replayEntity})
	if err != nil {
		if err.ErrCode() == bizerr.ErrCodeRecordNotFound {
			logger.Info("replay is not found, delete action completed")
			return true, nil
		}
		logger.Error("get replay status failed, try to destroy the replay", zap.Error(err))
		// try to destroy the replay
	}

	// execute the destroy operation
	if statusResp == nil || statusResp.Status == nil || statusResp.Phase != entity.ReplayPhaseDeleting {
		err := w.replayOperator.Destroy(ctx, &domain.DestroyReplayReq{Replay: replayEntity})
		if err != nil {
			logger.Error("destroy replay failed", zap.Error(err))
			return false, err.Chain("destroy replay")
		}
		return false, nil
	}

	// update the status information
	if statusResp.Status != nil {
		replayEntity.ReplayStatus = statusResp.Status
	}

	logger.Info("replay is deleting, waiting for the next sync")
	return false, nil
}

// ==================== Helper Methods ====================

// isDeployComplete check if the deploy is completed
func (w *replayWorkerImpl) isDeployComplete(resp *domain.GetReplayStatusResp) bool {
	if resp == nil || resp.Status == nil {
		return false
	}
	// can be determined based on the actual Replay resource status
	return resp.Phase == entity.ReplayPhaseRunning
}

// isRunning check if the replay is running
func (w *replayWorkerImpl) isRunning(resp *domain.GetReplayStatusResp) bool {
	if resp == nil || resp.Status == nil || resp.Phase != entity.ReplayPhaseRunning {
		return false
	}
	for _, task := range resp.Status.Tasks {
		for _, replica := range task.ReplicaStatuses {
			if replica.Phase != entity.ReplayTaskPhaseRunning {
				return false
			}
		}
	}
	return true
}

// isPaused check if the replay is paused
func (w *replayWorkerImpl) isPaused(resp *domain.GetReplayStatusResp) bool {
	if resp == nil || resp.Status == nil {
		return false
	}
	return resp.Phase == entity.ReplayPhaseStopped
}