package dataplane

import (
	"context"
	"fmt"

	"github.com/mitchellh/mapstructure"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/eks-provider/controllers/workflow/action"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

// restartOtherPod restart other pod besides the tidb cluster, like tiflow-executor pod
type restartOtherPod struct {
	cli            client.Client
	cachedProvider dataplane.EnvProviderManager
}

type restartOtherPodInputs struct {
	// ClusterNamespace is the namespace of the tidb cluster
	ClusterNamespace string `mapstructure:"clusterNamespace" json:"clusterNamespace" validate:"required"`
	// ClusterName is the tidb cluster name
	ClusterName string `mapstructure:"clusterName" json:"clusterName" validate:"required"`
	// Namespace is the namespace of the pod that needs to be restarted
	Namespace string `mapstructure:"namespace" json:"namespace" validate:"required"`
	// PodName is the name of the corresponding pod
	PodName string `mapstructure:"podName" json:"podName" validate:"required"`
	Force   bool   `mapstructure:"force" json:"force" validate:"required"`
}

var (
	_ action.Action = &restartOtherPod{}
)

const (
	componentLabelKey      = "app.kubernetes.io/component"
	gracefulRestartAnnoVal = "delete-pod"
)

var supportGracefulRestartAnnotations = map[string]string{}

func NewRestartOtherPod(cli client.Client, cachedProvider dataplane.EnvProviderManager) action.Action {
	return &restartOtherPod{cli: cli, cachedProvider: cachedProvider}
}

func (p *restartOtherPod) Name() cloudv1.WorkflowActionName {
	return consts.WfActionRestartOtherPodAction
}

func (p *restartOtherPod) PreCheck(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &restartOtherPodInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := validate.Struct(inputs); err != nil {
		return fmt.Errorf("validate inputs: %w", err)
	}
	return nil
}

func (p *restartOtherPod) Handle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	logger := log.FromContext(ctx)

	inputs := &restartOtherPodInputs{}
	_ = mapstructure.Decode(step.Inputs.Object, inputs)

	logger = logger.WithValues("clusterName", inputs.ClusterName, "podNamespace", inputs.Namespace, "podName", inputs.PodName)

	cluster := &cloudv1.Cluster{}
	if err := p.cli.Get(ctx, client.ObjectKey{Namespace: inputs.ClusterNamespace, Name: inputs.ClusterName}, cluster); err != nil {
		if apierrors.IsNotFound(err) {
			return cloudv1.NewWFNonRetriableError(fmt.Errorf("cluster %s not found", inputs.ClusterName))
		}
		return fmt.Errorf("get cluster: %w", err)
	}

	remoteCli, err := p.cachedProvider.GetClient(ctx, cluster)
	if err != nil {
		return fmt.Errorf("get dataplane cli: %w", err)
	}

	pod := &corev1.Pod{}
	if err := remoteCli.Get(ctx, client.ObjectKey{Namespace: inputs.Namespace, Name: inputs.PodName}, pod); err != nil {
		if apierrors.IsNotFound(err) {
			return cloudv1.NewWFNonRetriableError(fmt.Errorf("pod %s not found", inputs.PodName))
		}
		return fmt.Errorf("get pod: %w", err)
	}
	logger.Info("try to restart pod", "podName", pod.Name)
	if err := restartOnePod(ctx, remoteCli, pod, inputs.Force); err != nil {
		return fmt.Errorf("failed to restart pod %s: %w", pod.Name, err)
	}

	return nil
}

func restartOnePod(ctx context.Context, cli client.Client, pod *corev1.Pod, force bool) error {
	if pod.DeletionTimestamp != nil {
		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 := k8s.Apply(ctx, cli, pod); err != nil {
			return fmt.Errorf("failed to set restart annotation to pod %s/%s: %v", pod.Namespace, pod.Name, err)
		}
	} else {
		// delete the pod for restarting forcefully
		if err := cli.Delete(ctx, pod, client.GracePeriodSeconds(0)); err != nil {
			return fmt.Errorf("failed to delete pod %s/%s: %v", pod.Namespace, pod.Name, err)
		}
	}

	return nil
}