package hook

import (
	"context"

	"github.com/tidbcloud/pov-service/internal/common/errors"
	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/internal/pkg/log"
	"github.com/tidbcloud/pov-service/internal/service"
	"github.com/tidbcloud/pov-service/internal/service/workload/task"
)

type PreserveData struct {
	svc service.Workload
}

func NewPreserveData(svc service.Workload) PreserveData {
	return PreserveData{
		svc: svc,
	}
}

func (p PreserveData) Match(
	processor task.TaskProcessor,
	test model.WorkloadTest) bool {
	if processor.GetName() != task.PrepareDataRunningProcessorName {
		return false
	}

	isPrepared := test.IsDataPrepared()
	if !isPrepared {
		return false
	}
	if test.Cluster == nil {
		return false
	}

	// Check if there is valid preserved sample to enhance performance
	legacySample, err := p.svc.GetPreservedSampleDBs(
		context.Background(),
		test.WorkloadID,
	)
	if err != nil {
		if _, ok := err.(errors.InvalidArgument); !ok {
			return false
		}
	}
	if len(legacySample.Databases) > 0 && !legacySample.IsOutdated() {
		return false
	}

	return true
}

func (p PreserveData) Do(ctx context.Context, test model.WorkloadTest) error {
	log.Log.Sugar.Infof("Saving sample data for test %s", test.ID)

	return p.svc.PreserveWorkloadSample(ctx, test.WorkloadID, *test.Cluster)
}

func (p PreserveData) GetName() string {
	return "preserve_data_hook"
}