package commontasks

import (
	"context"
	"fmt"
	"strings"

	"github.com/pingcap/errors"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
)

type DrainNodeTask[T client.Object] struct {
	*reconciler.BaseTask[T]
	cm deploy.ClusterMeta
}

func NewDrainNodeTask[T client.Object](ctx *reconciler.Context[T], cm deploy.ClusterMeta) *DrainNodeTask[T] {
	return &DrainNodeTask[T]{
		BaseTask: &reconciler.BaseTask[T]{Context: ctx},
		cm:       cm,
	}
}

func (t *DrainNodeTask[T]) Name() string {
	return "DrainNode"
}

func (t *DrainNodeTask[T]) ClusterMeta() deploy.ClusterMeta {
	return t.cm
}

func (t *DrainNodeTask[T]) Create(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}

	annotations := t.Agent.GetAnnotations()
	drainValue := annotations[consts.AnnotationKeyO11yDrainNode]
	if drainValue == "" {
		return deploy.Done()
	}

	// Ensure that the annotation is removed after processing, to avoid repeated drain attempts
	defer func() {
		delete(annotations, consts.AnnotationKeyO11yDrainNode)
		t.Agent.SetAnnotations(annotations)
	}()

	parts := strings.Split(drainValue, "/")
	if len(parts) != 2 {
		return deploy.Error(fmt.Errorf("invalid drain-node annotation: %s", drainValue))
	}

	namespace := parts[0]
	podName := parts[1]
	if namespace == "" || podName == "" {
		return deploy.Error(fmt.Errorf("invalid drain-node annotation: %s, namespace/podName must not be empty", drainValue))
	}

	nodeName, err := t.drainNodePod(ctx, cli, namespace, podName)
	if err != nil {
		return deploy.Error(errors.Wrap(err, "failed to drain node"))
	}

	t.L.Info(fmt.Sprintf("Successfully initiated drain for node %s (pod: %s/%s)", nodeName, namespace, podName))
	return deploy.Done()
}

// drainNodePod applies the out-of-service taint to the node where the specified pod is running.
func (t *DrainNodeTask[T]) drainNodePod(ctx context.Context, cli client.Client, namespace, podName string) (string, error) {
	var pod corev1.Pod
	if err := cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: podName}, &pod); err != nil {
		return "", fmt.Errorf("pod %s/%s not found: %w", namespace, podName, err)
	}

	nodeName := pod.Spec.NodeName
	if nodeName == "" {
		return "", fmt.Errorf("pod %s/%s is not scheduled to any node", namespace, podName)
	}

	var node corev1.Node
	if err := cli.Get(ctx, client.ObjectKey{Name: nodeName}, &node); err != nil {
		return "", fmt.Errorf("node %s not found: %w", nodeName, err)
	}

	outOfServiceTaint := corev1.Taint{
		Key:    "node.kubernetes.io/out-of-service",
		Effect: corev1.TaintEffectNoExecute,
	}
	for _, taint := range node.Spec.Taints {
		if taint.Key == outOfServiceTaint.Key && taint.Effect == outOfServiceTaint.Effect {
			t.L.Info(fmt.Sprintf("Node %s already has out-of-service taint", nodeName))
			return nodeName, nil
		}
	}

	node.Spec.Taints = append(node.Spec.Taints, outOfServiceTaint)
	if err := cli.Update(ctx, &node); err != nil {
		return "", fmt.Errorf("failed to apply taint to node %s: %w", nodeName, err)
	}

	t.L.Info(fmt.Sprintf("Successfully applied out-of-service taint to node %s", nodeName))
	return nodeName, nil
}

func (t *DrainNodeTask[T]) Delete(ctx context.Context) deploy.Result {
	return deploy.Done()
}