package task

import (
	"context"
	"fmt"

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

const (
	prepareDataTriggerProcessorName = "prepare-data-trigger"
	PrepareDataRunningProcessorName = "prepare-data-running"
)

// PrepareDataTriggerProcessor handles triggering prepare data tasks
type PrepareDataTriggerProcessor struct {
	BaseProcessor
}

func NewPrepareDataTriggerProcessor(manager *ProcessorManager) *PrepareDataTriggerProcessor {
	return &PrepareDataTriggerProcessor{BaseProcessor{manager}}
}

func (p *PrepareDataTriggerProcessor) GetName() string {
	return prepareDataTriggerProcessorName
}

func (p *PrepareDataTriggerProcessor) CanProcess(test model.WorkloadTest, taskPhase model.WorkloadTaskPhase) bool {
	return test.Phase == model.WorkloadTestPhasePrepareData &&
		test.Status == model.ExecStatusRunning &&
		taskPhase == model.WorkloadTaskPhasePending
}

func (p *PrepareDataTriggerProcessor) Process(ctx context.Context, test model.WorkloadTest) (model.WorkloadTest, error) {
	log.Log.Infof("Triggering prepare data for test %s", test.ID)

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

	// Update task phase to triggering
	p.setTaskTriggering(currentTask)

	// Add testbed param
	test.AppendTestbedParam(test.Cluster.TestbedName)

	// Trigger prepare data - need to create params.RunWorkload
	runWorkloadParam := params.RunWorkload{
		WorkloadID:   test.WorkloadID,
		Name:         test.Name,
		WorkloadName: test.WorkloadName,
		Params:       &test.Params,
	}
	taskBase, err := p.manager.gateway.TriggerPrepareData(
		ctx, *test.Cluster, runWorkloadParam)
	if err != nil {
		// Don't return error for trigger processors to avoid retrying non-idempotent operations
		// Instead, mark the task as failed and let it be handled by phase transition processor
		log.Log.Errorf(err, "failed to trigger prepare data for test %s", test.ID)
		p.setTaskFailed(currentTask, err)
		return test, nil // Return nil to commit the transaction and avoid retry
	}

	// Update task with external service info
	p.setTaskRunning(currentTask, taskBase)

	return test, nil
}

// PrepareDataRunningProcessor handles running prepare data tasks
type PrepareDataRunningProcessor struct {
	BaseProcessor
}

func NewPrepareDataRunningProcessor(manager *ProcessorManager) *PrepareDataRunningProcessor {
	return &PrepareDataRunningProcessor{BaseProcessor{manager}}
}

func (p *PrepareDataRunningProcessor) GetName() string {
	return PrepareDataRunningProcessorName
}

func (p *PrepareDataRunningProcessor) CanProcess(test model.WorkloadTest, taskPhase model.WorkloadTaskPhase) bool {
	return test.Phase == model.WorkloadTestPhasePrepareData &&
		test.IsOnTheFly() &&
		taskPhase == model.WorkloadTaskPhaseRunning
}

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

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

	// Poll the task
	taskBase, result, err := p.manager.gateway.PollPrepareDataTask(
		ctx,
		*test.Cluster,
		currentTask.WorkloadTaskBase)
	if err != nil {
		log.Log.Errorf(err, "failed to poll prepare data task for test %s", test.ID)
		return test, nil // Don't fail, just continue polling
	}
	// TODO: get prepare progress by checking data size if possible

	// Update task with latest info
	p.updateTaskFromRunning(currentTask, taskBase)

	if result.TiWorkloadClusterID != "" {
		test.Cluster.TiWorkloadClusterID = result.TiWorkloadClusterID
	}
	if !test.Cluster.HasConnection() && result.HasConnection() {
		test.Cluster.Host = result.DBHost
		test.Cluster.Port = result.DBPort
	}

	// Check if completed
	switch taskBase.Status {
	case model.ExecStatusCompleted:
		currentTask.Result = result
		p.setTaskCompleted(currentTask)
		log.Log.Infof("Prepare data completed for test %s", test.ID)
	case model.ExecStatusFailed, model.ExecStatusStopped:
		p.setTaskFailedFromRunning(currentTask)
		log.Log.Errorf(fmt.Errorf("task failed"), "Prepare data failed for test %s: %s", test.ID, taskBase.Message)
	}

	return test, nil
}
