package common

import (
	"testing"

	"github.com/stretchr/testify/require"
	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"
)

func TestReplayTaskDelegatesPlacementToProvider(t *testing.T) {
	wantTask := cloudv1.ReplayTask{
		Name: "task-1",
		Resources: &cloudv1.ReplayTaskResources{
			InstanceType: "8C32G",
		},
	}
	wantAffinity := &corev1.Affinity{
		NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{
					{
						MatchExpressions: []corev1.NodeSelectorRequirement{
							{
								Key:      "node-spec",
								Operator: corev1.NodeSelectorOpIn,
								Values:   []string{"8C32G"},
							},
						},
					},
				},
			},
		},
	}
	wantTolerations := []corev1.Toleration{
		{
			Key:      "dedicated",
			Operator: corev1.TolerationOpEqual,
			Value:    "worker",
			Effect:   corev1.TaintEffectNoSchedule,
		},
	}
	provider := &replayProviderSpy{
		affinity:    wantAffinity,
		tolerations: wantTolerations,
	}
	replay := &replayTask{
		BaseTask: &task.BaseTask{
			Context: &task.Context{Provider: provider},
		},
	}

	gotAffinity := replay.buildReplayNodeAffinity(wantTask)
	gotTolerations := replay.buildReplayTolerations()

	require.Same(t, wantAffinity, gotAffinity)
	require.Equal(t, wantTolerations, gotTolerations)
	require.Equal(t, wantTask, provider.affinityTask)
}

func TestReplayTaskReturnsNilPlacementWithoutProvider(t *testing.T) {
	replay := &replayTask{
		BaseTask: &task.BaseTask{
			Context: &task.Context{},
		},
	}

	require.Nil(t, replay.buildReplayNodeAffinity(cloudv1.ReplayTask{}))
	require.Nil(t, replay.buildReplayTolerations())
}

type replayProviderSpy struct {
	affinity     *corev1.Affinity
	tolerations  []corev1.Toleration
	affinityTask cloudv1.ReplayTask
}

func (p *replayProviderSpy) Tasks() []deploy.Task { return nil }

func (p *replayProviderSpy) ReplayNodeAffinity(replayTask cloudv1.ReplayTask) *corev1.Affinity {
	p.affinityTask = replayTask
	return p.affinity
}

func (p *replayProviderSpy) ReplayTolerations() []corev1.Toleration {
	return p.tolerations
}

func (p *replayProviderSpy) ReplayZones() []string {
	return nil
}

func (p *replayProviderSpy) ReplaySubnets() []string {
	return nil
}