package eks

import (
	"context"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"

	"github.com/tidbcloud/eks-provider/controllers/network/task"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
)

const (
	cniContainerName     = "aws-node"
	cniInitContainerName = "aws-vpc-cni-init"
	// aws cni plugin default value is 1, which requires 6 or more IPs to be reserved per node
	// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/using-eni.html#AvailableIpPerENI
	// https://github.com/aws/amazon-vpc-cni-k8s#eni-allocation
	cniWarmENIs = "0"
	// aws cni plugin does not have default value for warm IP
	// for dedicated clusters, our IP allocation per node is nearly static except for control node
	// so a small warm up IPs is enough
	cniWarmIPs = "2"
)

const (
	// for both shared and dedicated
	cniEnvWarmENI = "WARM_ENI_TARGET"
	cniEnvWarmIP  = "WARM_IP_TARGET"
)

// CNI task deploys CNI plugin
// For EKS, the cni plugin is automatically deployed by eks itself
// but the default configuration has a limitation for 31 maximum nodes
// so we need to modify the default configuration according to
// https://docs.aws.amazon.com/eks/latest/userguide/cni-increase-ip-addresses.html
type CNI struct {
	*task.BaseTask
}

func NewCNI(ctx *task.Context) *CNI {
	return &CNI{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

func (t *CNI) Name() string {
	return "CNI"
}

func (t *CNI) Create(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	ds := &appsv1.DaemonSet{}
	key := types.NamespacedName{Namespace: "kube-system", Name: "aws-node"}
	if err := cli.Get(ctx, key, ds); err != nil {
		return deploy.Panic(err)
	}
	envVars := []corev1.EnvVar{
		{
			Name:  cniEnvWarmENI,
			Value: cniWarmENIs,
		},
		{
			Name:  cniEnvWarmIP,
			Value: cniWarmIPs,
		},
	}
	for _, envVar := range envVars {
		mutateDaemonSetEnv(ds, cniContainerName, envVar)
	}

	initEnvVars := []corev1.EnvVar{}
	for _, envVar := range initEnvVars {
		mutateDaemonSetInitEnv(ds, cniInitContainerName, envVar)
	}

	if err := cli.Update(ctx, ds); err != nil {
		return deploy.Panic(err)
	}

	return deploy.Done()
}

// mutateDaemonSetEnv mutates daemonset's container env var
func mutateDaemonSetEnv(ds *appsv1.DaemonSet, containerName string, envVar corev1.EnvVar) {
	mutateContainersSetEnv(ds.Spec.Template.Spec.Containers, containerName, envVar)
}

// mutateDaemonSetInitEnv mutates daemonset's init container env var
func mutateDaemonSetInitEnv(ds *appsv1.DaemonSet, containerName string, envVar corev1.EnvVar) {
	mutateContainersSetEnv(ds.Spec.Template.Spec.InitContainers, containerName, envVar)
}

func mutateContainersSetEnv(containers []corev1.Container, containerName string, envVar corev1.EnvVar) {
	for i := range containers {
		container := containers[i]
		if container.Name == containerName {
			exist := false
			for j := range container.Env {
				env := container.Env[j]
				if env.Name == envVar.Name {
					env.Value = envVar.Value
					exist = true
					container.Env[j] = env
				}
			}
			if !exist {
				container.Env = append(container.Env, envVar)
			}
			containers[i] = container
			return
		}
	}
}
