package service

import (
	"context"
	"fmt"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type PodRestartProvider interface {
	Restart(ctx context.Context, cli client.Client, pod *corev1.Pod, force bool) error
}

func NewPodRestarter(version string) PodRestartProvider {
	if version == "v2" {
		return &podRestarterV2{}
	}
	return &podRestarterV1{}
}

type podRestarterV1 struct{}

var _ PodRestartProvider = &podRestarterV1{}

func (p *podRestarterV1) Restart(ctx context.Context, cli client.Client, pod *corev1.Pod, force bool) error {
	if isDeleting(pod) {
		return nil
	}

	restartAnno, supportGracefulRestart := supportGracefulRestartAnnotations[pod.Labels[componentLabelKey]]
	if supportGracefulRestart && !force {
		// set the restart annotation to the pod for graceful restart
		if pod.Annotations == nil {
			pod.Annotations = make(map[string]string)
		}
		pod.Annotations[restartAnno] = gracefulRestartAnnoVal
		if err := cli.Update(ctx, pod); err != nil {
			return fmt.Errorf("failed to set restart annotation to pod %s/%s: %v", pod.Namespace, pod.Name, err)
		}
		return nil
	}

	// delete the pod for restarting forcefully
	return deletePod(ctx, cli, pod, []client.DeleteOption{client.GracePeriodSeconds(1)})
}

type podRestarterV2 struct{}

var _ PodRestartProvider = &podRestarterV2{}

func (p *podRestarterV2) Restart(ctx context.Context, cli client.Client, pod *corev1.Pod, force bool) error {
	if isDeleting(pod) {
		return nil
	}

	var (
		seconds int64
		opts    []client.DeleteOption
	)
	if pod.Labels[opv2.LabelKeyComponent] == opv2.LabelValComponentTiKV {
		// When delete a tikv pod, we have to set the grace period.
		seconds = tikvGracePeriodSeconds
	}
	if force {
		seconds = 1
	}
	if seconds > 0 {
		opts = append(opts, client.GracePeriodSeconds(seconds))
	}
	return deletePod(ctx, cli, pod, opts)
}

func deletePod(ctx context.Context, cli client.Client, pod *corev1.Pod, opts []client.DeleteOption) error {
	if err := cli.Delete(ctx, pod, opts...); err != nil {
		return fmt.Errorf("failed to delete pod %s/%s: %v", pod.Namespace, pod.Name, err)
	}
	return nil
}

func isDeleting(pod *corev1.Pod) bool {
	return pod.DeletionTimestamp != nil
}