package provider

import (
	corev1 "k8s.io/api/core/v1"

	"github.com/tidbcloud/eks-provider/controllers/replay/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
)

const (
	workerReuseNodeLabel = "node-spec"
	dedicatedTaintKey    = "dedicated"
	workerTaintValue     = "worker"
)

type premiumReplayProvider struct {
	tasks []deploy.Task
}

func newPremiumReplayProvider(tasks []deploy.Task) premiumReplayProvider {
	return premiumReplayProvider{tasks: tasks}
}

func (p *premiumReplayProvider) Tasks() []deploy.Task {
	return p.tasks
}

func (p *premiumReplayProvider) ReplayNodeAffinity(replayTask cloudv1.ReplayTask) *corev1.Affinity {
	instanceType := task.DefaultReplayNodeSpec
	if replayTask.Resources != nil && replayTask.Resources.InstanceType != "" {
		instanceType = replayTask.Resources.InstanceType
	}
	nodeSpecRequirement := corev1.NodeSelectorRequirement{
		Key:      workerReuseNodeLabel,
		Operator: corev1.NodeSelectorOpExists,
	}
	if instanceType != "" {
		nodeSpecRequirement.Operator = corev1.NodeSelectorOpIn
		nodeSpecRequirement.Values = []string{instanceType}
	}

	return &corev1.Affinity{
		NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{
					{
						MatchExpressions: []corev1.NodeSelectorRequirement{nodeSpecRequirement},
					},
				},
			},
		},
	}
}

func (p *premiumReplayProvider) ReplayTolerations() []corev1.Toleration {
	return []corev1.Toleration{
		{
			Key:      dedicatedTaintKey,
			Operator: corev1.TolerationOpEqual,
			Value:    workerTaintValue,
			Effect:   corev1.TaintEffectNoSchedule,
		},
	}
}