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"
)

type DedicatedEKSReplayProvider struct {
	replayName string
	tasks      []deploy.Task
	zones      []string
	subnets    []string
}

func NewDedicatedEKSReplayProvider(c *task.Context) *DedicatedEKSReplayProvider {
	if c == nil || c.Replay == nil {
		return nil
	}

	defaults := deriveAWSReplayDefaults(c.Replay)
	return &DedicatedEKSReplayProvider{
		tasks:      newDedicatedEKSReplayTasks(c),
		zones:      defaults.zones,
		subnets:    defaults.subnets,
		replayName: c.Replay.Name,
	}
}

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

func (p *DedicatedEKSReplayProvider) ReplayNodeAffinity(replayTask cloudv1.ReplayTask) *corev1.Affinity {
	instanceType := task.DefaultInstanceType
	if replayTask.Resources != nil && replayTask.Resources.InstanceType != "" {
		instanceType = replayTask.Resources.InstanceType
	}

	return &corev1.Affinity{
		NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{
					{
						MatchExpressions: []corev1.NodeSelectorRequirement{
							{
								Key:      task.NodeSelectorKey,
								Operator: corev1.NodeSelectorOpIn,
								Values:   []string{p.replayName},
							},
							{
								Key:      task.NodeInstanceTypeKey,
								Operator: corev1.NodeSelectorOpIn,
								Values:   []string{instanceType},
							},
						},
					},
				},
			},
		},
	}
}

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

func (p *DedicatedEKSReplayProvider) ReplayZones() []string {
	return p.zones
}

func (p *DedicatedEKSReplayProvider) ReplaySubnets() []string {
	return p.subnets
}