package service

import (
	"context"
	"fmt"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/samber/lo"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

const (
	tcName                   = "db"
	tcNamespacePrefix        = "tidb"
	componentLabelKey        = "app.kubernetes.io/component"
	InstanceLabelKey         = "app.kubernetes.io/instance"
	operationIDAnnotationKey = "cluster-agent.tidbcloud.com/operation-id"
	gracefulRestartAnnoVal   = "delete-pod"
)

var componentMap = map[string]struct{}{
	"pd":      {},
	"tikv":    {},
	"tidb":    {},
	"tiflash": {},
	"tiproxy": {},
}
var (
	clusterComponents []string
)

func init() {
	clusterComponents = lo.Keys(componentMap)
}

var supportGracefulRestartAnnotations = map[string]string{
	"pd":   v1alpha1.PDLeaderTransferAnnKey,     // https://github.com/pingcap/tidb-operator/pull/4896
	"tikv": v1alpha1.EvictLeaderAnnKey,          // https://docs.pingcap.com/tidb-in-kubernetes/stable/restart-a-tidb-cluster#perform-a-graceful-restart-to-a-single-tikv-pod
	"tidb": v1alpha1.TiDBGracefulShutdownAnnKey, // https://github.com/pingcap/tidb-operator/pull/4948
}

func (s *AgentService) Restart(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx).With("request", req)

	if err := validateRestartReq(req); err != nil {
		logger.Errorf("failed to validate the request: %v", err)
		return nil, err
	}

	restartFuncs := map[apiv1.RestartTargetType]func(context.Context, *apiv1.RestartReq) (*apiv1.RestartResp, error){
		apiv1.RestartTargetType_POD:       s.restartPod,
		apiv1.RestartTargetType_COMPONENT: s.restartComponent,
		apiv1.RestartTargetType_CLUSTER:   s.restartTiDBCluster,
		apiv1.RestartTargetType_NODE:      s.restartNode,
	}

	if restartFunc, ok := restartFuncs[req.TargetType]; ok {
		return restartFunc(ctx, req)
	}

	return &apiv1.RestartResp{}, nil
}

func validateRestartReq(req *apiv1.RestartReq) error {
	if req == nil {
		return fmt.Errorf("request is nil")
	}

	switch req.TargetType {
	case apiv1.RestartTargetType_POD, apiv1.RestartTargetType_NODE:
		if req.TargetName == "" {
			return fmt.Errorf("target name is required")
		}
	case apiv1.RestartTargetType_COMPONENT:
		if _, ok := componentMap[req.TargetName]; !ok {
			return fmt.Errorf("invalid component name: %s", req.TargetName)
		}
	}

	return nil
}

func (s *AgentService) restartPod(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx)
	ns := fmt.Sprintf("%s%s", tcNamespacePrefix, req.ClusterId)

	var pod v1.Pod
	if err := s.dataplaneCli.Get(ctx, client.ObjectKey{
		Name:      req.TargetName,
		Namespace: ns,
	}, &pod); err != nil {
		msg := fmt.Sprintf("failed to get pod %s: %v", req.TargetName, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}

	if err := s.podRestarter.Restart(ctx, s.dataplaneCli, &pod, req.Force); err != nil {
		msg := fmt.Sprintf("failed to restart pod %s/%s: %v", pod.Namespace, pod.Name, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}

	return &apiv1.RestartResp{Message: "success"}, nil
}

func (s *AgentService) restartComponent(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx)

	// Ignore to check the tidbcluster's status when using force restart
	if req.Force {
		// restart all pods of the component forcefully by default
		var matchLabels = client.MatchingLabels{componentLabelKey: req.TargetName}
		// if target group is specified, only restart the pods in the group
		if req.TargetGroup != "" {
			matchLabels[InstanceLabelKey] = req.TargetGroup
		}
		var compPods v1.PodList
		if err := s.dataplaneCli.List(ctx, &compPods, client.InNamespace(buildNamespace(req.ClusterId)), matchLabels); err != nil {
			msg := fmt.Sprintf("failed to list pods for component %s: %v", req.TargetName, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: msg}, err
		}

		for _, pod := range compPods.Items {
			if err := s.podRestarter.Restart(ctx, s.dataplaneCli, &pod, req.Force); err != nil {
				msg := fmt.Sprintf("failed to restart pod %s/%s: %v", pod.Namespace, pod.Name, err)
				logger.Error(msg)
				return &apiv1.RestartResp{Message: msg}, err
			}
		}
		return &apiv1.RestartResp{Message: "success"}, nil
	} else {
		if s.UseTiDBOperatorV2() {
			return s.graceRestartAllGroupsOfComponent4V2(ctx, req)
		}

		return s.graceRestartAllGroupsOfComponent4V1(ctx, req)
	}
}

func (s *AgentService) graceRestartAllGroupsOfComponent4V1(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx).With("func", "restartAllGroupsOfComponent4V1")

	name2tc, err := s.getTCsByClusterId(ctx, req.ClusterId)
	if err != nil {
		return nil, err
	}

	if req.TargetGroup != "" {
		tc, ok := name2tc[req.TargetGroup]
		if !ok {
			return nil, status.Errorf(codes.NotFound, "tidbcluster %s not found in cluster %s", req.TargetGroup, req.ClusterId)
		}
		name2tc = map[string]*v1alpha1.TidbCluster{
			req.TargetGroup: tc,
		}
	}
	if len(name2tc) == 0 {
		logger.Info("empty tidbcluster list, nothing to do")
	}
	logger.Infof("try to restart component %s for tcs: %s", req.TargetName, lo.Keys(name2tc))
	for _, tc := range name2tc {
		if err := s.graceRestartComponentGroup4V1(ctx, tc, req.TargetName, req.OperationId); err != nil {
			return nil, err
		}
	}
	return &apiv1.RestartResp{Message: "success"}, nil
}

// graceRestartComponentGroup4V1 handles the restart of a component group for TiDB Operator v1.
func (s *AgentService) graceRestartComponentGroup4V1(ctx context.Context, tc *v1alpha1.TidbCluster, component string, operationId string) error {
	logger := tilog.WithContextS(ctx).With("func", "restartComponentGroup4V1")

	// Use annotations to trigger the graceful restart of the component pods
	changed := setAnnotationForComponent(tc, component, operationIDAnnotationKey, operationId)
	if !changed {
		return nil
	}
	// TODO(ideascf): do we really need to check tc status here?
	if !isTCReady(tc) {
		return status.Errorf(codes.FailedPrecondition, "tc(%s) is not ready", tc.Name)
	}
	if err := k8s.Apply(ctx, s.dataplaneCli, tc); err != nil {
		msg := fmt.Sprintf("failed to set annotations for %s in tidbcluster CR(%s): %v", component, tc.Name, err)
		logger.Error(msg)
		return err
	}
	return nil
}

func (s *AgentService) graceRestartAllGroupsOfComponent4V2(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx).With("func", "restartAllGroupsOfComponent4V2")
	ns := buildNamespace(req.ClusterId)

	var groupList client.ObjectList
	switch req.TargetName {
	case string(cloudv1.ComponentPD):
		groupList = &opv2.PDGroupList{}
	case string(cloudv1.ComponentTiProxy):
		groupList = &opv2.TiProxyGroupList{}
	case string(cloudv1.ComponentTiDB):
		groupList = &opv2.TiDBGroupList{}
	case string(cloudv1.ComponentTiKV):
		groupList = &opv2.TiKVGroupList{}
	case string(cloudv1.ComponentTiFlash):
		groupList = &opv2.TiFlashGroupList{}
	default:
		err := fmt.Errorf("unsupported component type: %s", req.TargetName)
		return &apiv1.RestartResp{Message: err.Error()}, err
	}

	if err := s.dataplaneCli.List(ctx, groupList, client.InNamespace(ns)); err != nil {
		msg := fmt.Sprintf("failed to list groups for component %s: %v", req.TargetName, err)
		return &apiv1.RestartResp{Message: msg}, err
	}

	targets := []client.Object{}
	switch list := groupList.(type) {
	case *opv2.PDGroupList:
		for i := range list.Items {
			targets = append(targets, &list.Items[i])
		}
	case *opv2.TiProxyGroupList:
		for i := range list.Items {
			targets = append(targets, &list.Items[i])
		}
	case *opv2.TiDBGroupList:
		for i := range list.Items {
			targets = append(targets, &list.Items[i])
		}
	case *opv2.TiKVGroupList:
		for i := range list.Items {
			targets = append(targets, &list.Items[i])
		}
	case *opv2.TiFlashGroupList:
		for i := range list.Items {
			targets = append(targets, &list.Items[i])
		}
	}
	if req.TargetGroup != "" {
		targets = lo.Filter(targets, func(target client.Object, _ int) bool {
			return target.GetName() == req.TargetGroup
		})
		if len(targets) == 0 {
			return nil, status.Errorf(codes.NotFound, "target group %s not found for component %s", req.TargetGroup, req.TargetName)
		}
	}
	if len(targets) == 0 {
		logger.Info("empty tidbcluster list, nothing to do")
	}
	names := lo.Map(targets, func(target client.Object, _ int) string {
		return target.GetName()
	})
	logger.Infof("try to restart component %s for groups: %v", req.TargetName, names)
	for _, target := range targets {
		if err := s.graceRestartComponentGroup4V2(ctx, target, req.OperationId); err != nil {
			return nil, err
		}
	}

	return &apiv1.RestartResp{Message: "success"}, nil
}

// graceRestartComponentGroup4V2 handles the restart of a component group for TiDB Operator v2.
func (s *AgentService) graceRestartComponentGroup4V2(ctx context.Context, target client.Object, operationId string) error {
	logger := tilog.WithContextS(ctx).With("func", "restartComponentGroup", "operationId", operationId)

	var (
		conditions []metav1.Condition
	)

	switch target := target.(type) {
	case *opv2.PDGroup:
		conditions = target.Status.Conditions
		if target.Spec.Template.Annotations[operationIDAnnotationKey] != operationId {
			if target.Spec.Template.Annotations == nil {
				target.Spec.Template.Annotations = make(map[string]string)
			}
			target.Spec.Template.Annotations[operationIDAnnotationKey] = operationId
		}
	case *opv2.TiProxyGroup:
		conditions = target.Status.Conditions
		if target.Spec.Template.Annotations[operationIDAnnotationKey] != operationId {
			if target.Spec.Template.Annotations == nil {
				target.Spec.Template.Annotations = make(map[string]string)
			}
			target.Spec.Template.Annotations[operationIDAnnotationKey] = operationId
		}
	case *opv2.TiDBGroup:
		conditions = target.Status.Conditions
		if target.Spec.Template.Annotations[operationIDAnnotationKey] != operationId {
			if target.Spec.Template.Annotations == nil {
				target.Spec.Template.Annotations = make(map[string]string)
			}
			target.Spec.Template.Annotations[operationIDAnnotationKey] = operationId
		}
	case *opv2.TiKVGroup:
		conditions = target.Status.Conditions
		if target.Spec.Template.Annotations[operationIDAnnotationKey] != operationId {
			if target.Spec.Template.Annotations == nil {
				target.Spec.Template.Annotations = make(map[string]string)
			}
			target.Spec.Template.Annotations[operationIDAnnotationKey] = operationId
		}
	case *opv2.TiFlashGroup:
		conditions = target.Status.Conditions
		if target.Spec.Template.Annotations[operationIDAnnotationKey] != operationId {
			if target.Spec.Template.Annotations == nil {
				target.Spec.Template.Annotations = make(map[string]string)
			}
			target.Spec.Template.Annotations[operationIDAnnotationKey] = operationId
		}
	default:
		return status.Errorf(codes.Internal, "unsupported component: %T(%s)", target, target.GetName())
	}

	ready, synced := meta.IsStatusConditionTrue(conditions, opv2.CondReady), meta.IsStatusConditionTrue(conditions, opv2.CondSynced)
	if !ready || !synced {
		err := fmt.Errorf("group is not ready or synced, ready: %v, synced: %v", ready, synced)
		logger.Error(err.Error())
		return status.Errorf(codes.FailedPrecondition, "group is not ready or synced, ready: %v, synced: %v", ready, synced)
	}

	if err := s.dataplaneCli.Update(ctx, target); err != nil {
		msg := fmt.Sprintf("failed to update the group: %v", err)
		logger.Error(msg)
		return err
	}
	return nil
}

func (s *AgentService) restartTiDBCluster(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx)

	// Ignore to check the tidbcluster's status when using force restart
	if req.Force {
		// restart all pods of the tidb cluster forcefully
		var pods v1.PodList
		if err := s.dataplaneCli.List(
			ctx,
			&pods,
			client.InNamespace(buildNamespace(req.ClusterId)),
			client.HasLabels{componentLabelKey},
		); err != nil {
			msg := fmt.Sprintf("failed to list pods for tidbcluster: %v", err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: msg}, err
		}

		for _, pod := range pods.Items {
			if _, ok := componentMap[pod.Labels[componentLabelKey]]; ok {
				if err := s.podRestarter.Restart(ctx, s.dataplaneCli, &pod, req.Force); err != nil {
					msg := fmt.Sprintf("failed to restart pod %s/%s: %v", pod.Namespace, pod.Name, err)
					logger.Error(msg)
					return &apiv1.RestartResp{Message: msg}, err
				}
			}
		}
	} else {
		if s.UseTiDBOperatorV2() {
			err := fmt.Errorf("graceful restart the whole tidb cluster is not supported in tidb operator v2, restart component groups instead")
			logger.Error(err.Error())
			return &apiv1.RestartResp{Message: err.Error()}, err
		}

		tc, err := s.getAndCheckTiDBCluster(ctx, req.ClusterId, req.TargetGroup)
		if err != nil {
			logger.Error(err.Error())
			return &apiv1.RestartResp{Message: err.Error()}, err
		}

		for _, comp := range tc.AllComponentSpec() {
			setAnnotationForComponent(tc, comp.MemberType().String(), operationIDAnnotationKey, req.OperationId)
		}
		if err := k8s.Apply(ctx, s.dataplaneCli, tc); err != nil {
			msg := fmt.Sprintf("failed to set annotations for %s in tidbcluster CR: %v", req.TargetName, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: msg}, err
		}
	}

	return &apiv1.RestartResp{Message: "success"}, nil
}

func (s *AgentService) restartNode(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx)

	// Get the pod and check if it's a tidb component pod
	ns := fmt.Sprintf("%s%s", tcNamespacePrefix, req.ClusterId)
	var pod v1.Pod
	if err := s.dataplaneCli.Get(ctx, client.ObjectKey{
		Name:      req.TargetName,
		Namespace: ns,
	}, &pod); err != nil {
		msg := fmt.Sprintf("failed to get pod %s: %v", req.TargetName, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}
	if _, ok := componentMap[pod.Labels[componentLabelKey]]; !ok {
		err := fmt.Errorf("can not restart a non-tidb-component node, pod: %s/%s", ns, req.TargetName)
		logger.Error(err)
		return &apiv1.RestartResp{Message: err.Error()}, err
	}

	// Get the node and check if it has been restarted by the same operation id
	var node v1.Node
	if err := s.dataplaneCli.Get(ctx, client.ObjectKey{Name: pod.Spec.NodeName}, &node); err != nil {
		msg := fmt.Sprintf("failed to get node %s of pod %s/%s: %v", pod.Spec.NodeName, ns, req.TargetName, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}
	val, ok := node.GetAnnotations()[operationIDAnnotationKey]
	if ok && val == req.OperationId {
		logger.Infof("node %s has been restarted by operation %s", node.Name, req.OperationId)
		return &apiv1.RestartResp{Message: "success"}, nil
	}

	if err := s.nodeRestarter.Restart(ctx, node.Spec.ProviderID); err != nil {
		msg := fmt.Sprintf("failed to reboot node %s: %v", node.Spec.ProviderID, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}

	// set the operation id to the node
	if node.GetAnnotations() == nil {
		node.Annotations = make(map[string]string)
	}
	node.Annotations[operationIDAnnotationKey] = req.OperationId
	if err := k8s.Apply(ctx, s.dataplaneCli, &node); err != nil {
		msg := fmt.Sprintf("failed to set annotations for node %s: %v", node.Name, err)
		logger.Error(msg)
		return &apiv1.RestartResp{Message: msg}, err
	}

	return &apiv1.RestartResp{Message: "success"}, nil
}

func (s *AgentService) getAndCheckTiDBCluster(ctx context.Context, clusterId, targetTC string) (*v1alpha1.TidbCluster, error) {
	tc, err := s.getTiDBClusterV1(ctx, clusterId, targetTC)
	if err != nil {
		return nil, err
	}
	ready := isTCReady(tc)
	if !ready {
		return nil, fmt.Errorf("cluster %s's tc is not ready", clusterId)
	}
	return tc, nil
}

func (s *AgentService) getTiDBClusterV1(ctx context.Context, clusterId, targetTC string) (*v1alpha1.TidbCluster, error) {
	ns := buildNamespace(clusterId)

	name := tcName
	if targetTC != "" {
		name = targetTC
	}

	tc := &v1alpha1.TidbCluster{}
	if err := s.dataplaneCli.Get(ctx, client.ObjectKey{Namespace: ns, Name: name}, tc); err != nil {
		return nil, fmt.Errorf("failed to get tidbcluster CR: %w", err)
	}

	return tc, nil
}