package task

import (
	"context"
	"fmt"

	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/internal/pkg/log"
)

const (
	StopTestRetryCount = 3
)

// StopTestProcessor handles stopping running tests
type StopTestProcessor struct {
	BaseProcessor
	// It's chance that the task has finished when processor try to stop it.
	// Retry to handle this case
	retryCount        int
	testStopFailCount map[string]int
}

func NewStopTestProcessor(manager *ProcessorManager) *StopTestProcessor {
	return &StopTestProcessor{
		BaseProcessor:     BaseProcessor{manager},
		retryCount:        StopTestRetryCount,
		testStopFailCount: make(map[string]int),
	}
}

func (p *StopTestProcessor) GetName() string {
	return "stop-test"
}

func (p *StopTestProcessor) CanProcess(test model.WorkloadTest, taskPhase model.WorkloadTaskPhase) bool {
	return test.Status == model.ExecStatusStopping &&
		// Running task can be stopped by calling external service
		(taskPhase == model.WorkloadTaskPhaseRunning ||
			// while pending task can be set to stop directly
			taskPhase == model.WorkloadTaskPhasePending)
}

func (p *StopTestProcessor) Process(ctx context.Context, test model.WorkloadTest) (model.WorkloadTest, error) {
	log.Log.Infof("Stopping test %s", test.ID)

	currentTask := test.CurrentPhaseTask()
	if currentTask == nil {
		return test, fmt.Errorf("no current task found")
	}

	// Pending task can be set to stop directly
	if currentTask.Phase == model.WorkloadTaskPhasePending {
		p.setTaskStopped(currentTask)
		return test, nil
	}

	// Update task phase to stopping
	p.setTaskStopping(currentTask)

	err := p.manager.gateway.CancelTask(ctx, currentTask.WorkloadTaskBase)
	if err != nil {
		failedCount := p.increaseFailedCount(test.ID)
		if failedCount > p.retryCount {
			// If failedCount > retryCount, set task phase to failed
			p.setTaskFailed(currentTask, err)
			// Remove failed count to avoid mem leak
			p.removeFailedCount(test.ID)
		} else {
			// else, reset task phase to running for next retry
			p.setTaskRunning(currentTask, currentTask.WorkloadTaskBase)
		}
	}

	return test, nil
}

func (p *StopTestProcessor) increaseFailedCount(testID string) int {
	if v, ok := p.testStopFailCount[testID]; ok {
		p.testStopFailCount[testID] = v + 1
		return v + 1
	}
	p.testStopFailCount[testID] = 1
	return 1
}

func (p *StopTestProcessor) removeFailedCount(testID string) {
	delete(p.testStopFailCount, testID)
}

// CreateClusterRunningProcessor handles running create cluster tasks
type StopTestRunningProcessor struct {
	BaseProcessor
}

func NewStopTestRunningProcessor(manager *ProcessorManager) *StopTestRunningProcessor {
	return &StopTestRunningProcessor{BaseProcessor{manager}}
}

func (p *StopTestRunningProcessor) GetName() string {
	return "stop-test-running"
}

func (p *StopTestRunningProcessor) CanProcess(test model.WorkloadTest, taskPhase model.WorkloadTaskPhase) bool {
	return test.Status == model.ExecStatusStopping &&
		taskPhase == model.WorkloadTaskPhaseStopping
}

func (p *StopTestRunningProcessor) Process(ctx context.Context, test model.WorkloadTest) (model.WorkloadTest, error) {
	log.Log.Infof("Polling stopping tasks for test %s", test.ID)

	currentTask := test.CurrentPhaseTask()
	if currentTask == nil {
		return test, fmt.Errorf("no current task found")
	}

	var taskBase model.WorkloadTaskBase
	switch currentTask.Typ {
	case model.WorkloadTaskTypeCreateCluster:
		if base, _, err := p.manager.gateway.PollCreateClusterTask(
			ctx,
			test.Cluster,
			currentTask.WorkloadTaskBase); err != nil {
			log.Log.Errorf(err, "failed to poll create cluster task for test %s", test.ID)
			return test, nil // Don't fail, just continue polling
		} else {
			taskBase = base
		}
	case model.WorkloadTaskTypePrepareData:
		if base, _, err := p.manager.gateway.PollPrepareDataTask(
			ctx,
			*test.Cluster,
			currentTask.WorkloadTaskBase); err != nil {
			log.Log.Errorf(err, "failed to poll prepare data to poll  task for test %s", test.ID)
			return test, nil // Don't fail, just continue polling
		} else {
			taskBase = base
		}
	case model.WorkloadTaskTypeRunTraffic:
		if base, err := p.manager.gateway.PollRunTrafficTask(ctx, currentTask.WorkloadTaskBase); err != nil {
			log.Log.Errorf(err, "failed to poll run traffic task for test %s", test.ID)
			return test, nil // Don't fail, just continue polling
		} else {
			taskBase = base
		}
	default:
		return test, fmt.Errorf("unknown task type: %s", currentTask.Typ)
	}

	if taskBase.Status == model.ExecStatusStopped {
		p.setTaskStopped(currentTask)
		log.Log.Infof("Task %s stopped for test %s", currentTask.Typ, test.ID)
	}

	return test, nil
}