package commontasks

import (
	"context"
	"fmt"

	"github.com/hashicorp/go-multierror"
	"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 RestartTask[T client.Object] struct {
	*reconciler.BaseTask[T]

	cm            deploy.ClusterMeta
	agentType     string
	namespace     string
	enableVector  bool
	enableVMAgent bool
}

func NewVMAgentRestartTask[T client.Object](
	ctx *reconciler.Context[T],
	cm deploy.ClusterMeta,
	agentType string,
	namespace string,
) *RestartTask[T] {
	return newRestartTask(ctx, cm, agentType, namespace, false, true)
}

func NewVectorRestartTask[T client.Object](
	ctx *reconciler.Context[T],
	cm deploy.ClusterMeta,
	agentType string,
	namespace string,
) *RestartTask[T] {
	return newRestartTask(ctx, cm, agentType, namespace, true, false)
}

func NewVectorVMAgentRestartTask[T client.Object](
	ctx *reconciler.Context[T],
	cm deploy.ClusterMeta,
	agentType string,
	namespace string,
) *RestartTask[T] {
	return newRestartTask(ctx, cm, agentType, namespace, true, true)
}

func newRestartTask[T client.Object](
	ctx *reconciler.Context[T],
	cm deploy.ClusterMeta,
	agentType string,
	namespace string,
	enableVector bool,
	enableVMAgent bool,
) *RestartTask[T] {
	return &RestartTask[T]{
		BaseTask:      &reconciler.BaseTask[T]{Context: ctx},
		cm:            cm,
		agentType:     agentType,
		namespace:     namespace,
		enableVector:  enableVector,
		enableVMAgent: enableVMAgent,
	}
}

func (t *RestartTask[T]) Name() string {
	return "Restart"
}

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

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

	annotations := t.Agent.GetAnnotations()
	if annotations[consts.AnnotationKeyO11yVectorRestart] != "true" && annotations[consts.AnnotationKeyO11yVMAgentRestart] != "true" {
		return deploy.Done()
	}

	modified, err := t.restartComponents(ctx, cli, annotations)
	if modified {
		t.Agent.SetAnnotations(annotations)
	}

	if err != nil {
		return deploy.Panic(err)
	}
	return deploy.Done()
}

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

func (t *RestartTask[T]) restartComponents(ctx context.Context, cli client.Client, annotations map[string]string) (bool, error) {
	var modified bool
	var errs error

	if t.enableVector && annotations[consts.AnnotationKeyO11yVectorRestart] == "true" {
		delete(annotations, consts.AnnotationKeyO11yVectorRestart)
		modified = true
		if err := t.restartVectorPods(ctx, cli); err != nil {
			errs = multierror.Append(errs, err)
		}
	}

	if t.enableVMAgent && annotations[consts.AnnotationKeyO11yVMAgentRestart] == "true" {
		delete(annotations, consts.AnnotationKeyO11yVMAgentRestart)
		modified = true
		if err := t.restartVMAgentPods(ctx, cli); err != nil {
			errs = multierror.Append(errs, err)
		}
	}

	return modified, errs
}

func (t *RestartTask[T]) restartVectorPods(ctx context.Context, cli client.Client) error {
	return t.restartPods(ctx, cli, consts.VectorLabels(t.agentType), consts.ComponentVector)
}

func (t *RestartTask[T]) restartVMAgentPods(ctx context.Context, cli client.Client) error {
	return t.restartPods(ctx, cli, consts.VMAgentLabels(t.agentType), consts.ComponentVMAgent)
}

func (t *RestartTask[T]) restartPods(ctx context.Context, cli client.Client, labels map[string]string, component string) error {
	t.L.Info("Restarting pods", "component", component, "labels", labels, "namespace", t.namespace)

	var podList corev1.PodList
	if err := cli.List(ctx, &podList, client.InNamespace(t.namespace), client.MatchingLabels(labels)); err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to list %s pods", component))
	}

	for i := 0; i < len(podList.Items); i++ {
		pod := podList.Items[i]
		if pod.DeletionTimestamp != nil {
			continue
		}
		if err := cli.Delete(ctx, &pod); err != nil {
			return errors.Wrapf(err, "failed to delete pod %s/%s", pod.Namespace, pod.Name)
		}
		t.L.Info(fmt.Sprintf("Deleted %s pod: %s/%s", component, pod.Namespace, pod.Name))
	}
	return nil
}