package helper

import (
	"context"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/util/retry"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/observability/pkg/diagcontroller/config"
)

const (
	OfflineAnnotationKey = "diag-engine/offline"
)

func ListDiagEnginePods(cli client.Client, conf config.DiagEngineConfig) (*corev1.PodList, error) {
	pods := &corev1.PodList{}
	opts := &client.ListOptions{
		Namespace: conf.Namespace,
		LabelSelector: labels.SelectorFromSet(labels.Set{
			"app": conf.Name,
		}),
	}

	err := cli.List(context.Background(), pods, opts)
	if err != nil {
		return nil, err
	}
	return pods, nil
}

func ListAvailableEnginePods(cli client.Client, conf config.DiagEngineConfig) (*corev1.PodList, error) {
	pods, err := ListDiagEnginePods(cli, conf)
	if err != nil {
		return nil, err
	}
	available := &corev1.PodList{}
	for _, pod := range pods.Items {
		if pod.Annotations == nil {
			available.Items = append(available.Items, pod)
			continue
		}
		if key, ok := pod.Annotations[OfflineAnnotationKey]; !ok || key != "true" {
			available.Items = append(available.Items, pod)
		}
	}
	return available, nil
}

func ListOfflineDiagEnginePods(cli client.Client, conf config.DiagEngineConfig) (*corev1.PodList, error) {
	offline := &corev1.PodList{}
	pods, err := ListDiagEnginePods(cli, conf)
	if err != nil {
		return offline, err
	}

	for _, pod := range pods.Items {
		if pod.Annotations == nil {
			continue
		}
		if key, ok := pod.Annotations[OfflineAnnotationKey]; ok && key == "true" {
			offline.Items = append(offline.Items, pod)
		}
	}

	return offline, nil
}

func MarkDiagEngineOffline(cli client.Client, engine *corev1.Pod) error {
	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		if err := cli.Get(context.Background(), types.NamespacedName{Namespace: engine.Namespace, Name: engine.Name}, engine); err != nil {
			return err
		}
		if engine.Annotations == nil {
			engine.Annotations = map[string]string{}
		}
		engine.Annotations[OfflineAnnotationKey] = "true"
		if err := cli.Update(context.Background(), engine); err != nil {
			return err
		}
		return nil
	})
}