package service

import (
	"context"
	"fmt"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/samber/lo"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	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"
)

var componentMap = map[string]struct{}{
	PDComponentType:                {},
	TSOComponentType:               {},
	SchedulingComponentType:        {},
	TiKVComponentType:              {},
	TiDBComponentType:              {},
	TiDBWorkerComponentType:        {},
	TiProxyComponentType:           {},
	TiFlashComputeComponentType:    {},
	TiFlashWriteComponentType:      {},
	TiKVWorkerComponentType:        {},
	CoprocessorWorkerComponentType: {},
}
var (
	clusterComponents []string
)

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

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 req.TargetName == "" {
			return fmt.Errorf("target name is required")
		}
		if req.TargetGroup == "" {
			return fmt.Errorf("target group is required for component restart")
		}

		compType := convertComponentTypeFromAPI(req.TargetName)
		if _, ok := componentMap[compType]; !ok {
			return fmt.Errorf("invalid component type: %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", namespacePrefix, 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{K8sComponentLabelKey: req.TargetName}
		// if target group is specified, only restart the pods in the group
		matchLabels[K8sInstanceLabelKey] = 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 %s: %v", req.TargetName, req.TargetGroup, 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
	}

	return s.restartComponentGroup(ctx, req)
}

func (s *AgentService) restartComponentGroup(ctx context.Context, req *apiv1.RestartReq) (*apiv1.RestartResp, error) {
	logger := tilog.WithContextS(ctx).With("func", "restartComponentGroup", "req", req)
	ns := buildNamespace(req.ClusterId)
	key := client.ObjectKey{Namespace: ns, Name: req.TargetGroup}

	var (
		conditions []metav1.Condition
		group      client.Object
	)

	switch convertComponentTypeFromAPI(req.TargetName) {
	case PDComponentType:
		var pdg opv2.PDGroup
		if err := s.dataplaneCli.Get(ctx, key, &pdg); err != nil {
			msg := fmt.Sprintf("failed to get the pd group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = pdg.Status.Conditions
		if pdg.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if pdg.Spec.Template.Annotations == nil {
				pdg.Spec.Template.Annotations = make(map[string]string)
			}
			pdg.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &pdg
		}
	case TiDBComponentType:
		var dbg opv2.TiDBGroup
		if err := s.dataplaneCli.Get(ctx, key, &dbg); err != nil {
			msg := fmt.Sprintf("failed to get the tidb group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = dbg.Status.Conditions
		if dbg.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if dbg.Spec.Template.Annotations == nil {
				dbg.Spec.Template.Annotations = make(map[string]string)
			}
			dbg.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &dbg
		}
	case TiKVComponentType:
		var kvg opv2.TiKVGroup
		if err := s.dataplaneCli.Get(ctx, key, &kvg); err != nil {
			msg := fmt.Sprintf("failed to get the tikv group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = kvg.Status.Conditions
		if kvg.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if kvg.Spec.Template.Annotations == nil {
				kvg.Spec.Template.Annotations = make(map[string]string)
			}
			kvg.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &kvg
		}
	case TiFlashComputeComponentType:
		var flash opv2.TiFlashGroup
		key.Name = fmt.Sprintf("%s-compute", req.TargetGroup)
		if err := s.dataplaneCli.Get(ctx, key, &flash); err != nil {
			msg := fmt.Sprintf("failed to get the tiflash compute group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = flash.Status.Conditions
		if flash.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if flash.Spec.Template.Annotations == nil {
				flash.Spec.Template.Annotations = make(map[string]string)
			}
			flash.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &flash
		}
	case TiFlashWriteComponentType:
		var flash opv2.TiFlashGroup
		key.Name = fmt.Sprintf("%s-write", req.TargetGroup)
		if err := s.dataplaneCli.Get(ctx, key, &flash); err != nil {
			msg := fmt.Sprintf("failed to get the tiflash write group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = flash.Status.Conditions
		if flash.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if flash.Spec.Template.Annotations == nil {
				flash.Spec.Template.Annotations = make(map[string]string)
			}
			flash.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &flash
		}
	case TiDBWorkerComponentType:
		var worker opv2.TiDBGroup
		key.Name = fmt.Sprintf("%s-worker", req.TargetGroup)
		if err := s.dataplaneCli.Get(ctx, key, &worker); err != nil {
			msg := fmt.Sprintf("failed to get the tidb worker group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = worker.Status.Conditions
		if worker.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if worker.Spec.Template.Annotations == nil {
				worker.Spec.Template.Annotations = make(map[string]string)
			}
			worker.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &worker
		}
	case TiProxyComponentType:
		var proxy opv2.TiProxyGroup
		if err := s.dataplaneCli.Get(ctx, key, &proxy); err != nil {
			msg := fmt.Sprintf("failed to get the tiproxy group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = proxy.Status.Conditions
		if proxy.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if proxy.Spec.Template.Annotations == nil {
				proxy.Spec.Template.Annotations = make(map[string]string)
			}
			proxy.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &proxy
		}
	case TSOComponentType:
		var tso opv2.TSOGroup
		if err := s.dataplaneCli.Get(ctx, key, &tso); err != nil {
			msg := fmt.Sprintf("failed to get the pd-tso group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = tso.Status.Conditions
		if tso.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if tso.Spec.Template.Annotations == nil {
				tso.Spec.Template.Annotations = make(map[string]string)
			}
			tso.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &tso
		}
	case SchedulingComponentType:
		var scheduling opv2.SchedulingGroup
		if err := s.dataplaneCli.Get(ctx, key, &scheduling); err != nil {
			msg := fmt.Sprintf("failed to get the pd-scheduling group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = scheduling.Status.Conditions
		if scheduling.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if scheduling.Spec.Template.Annotations == nil {
				scheduling.Spec.Template.Annotations = make(map[string]string)
			}
			scheduling.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &scheduling
		}
	case TiKVWorkerComponentType:
		var worker opv2.TiKVWorkerGroup
		key.Name = fmt.Sprintf("%s-default", req.TargetGroup)
		if err := s.dataplaneCli.Get(ctx, key, &worker); err != nil {
			msg := fmt.Sprintf("failed to get the tikv worker group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = worker.Status.Conditions
		if worker.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if worker.Spec.Template.Annotations == nil {
				worker.Spec.Template.Annotations = make(map[string]string)
			}
			worker.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &worker
		}
	case CoprocessorWorkerComponentType:
		var worker opv2.TiKVWorkerGroup
		key.Name = fmt.Sprintf("%s-coprocessor", req.TargetGroup)
		if err := s.dataplaneCli.Get(ctx, key, &worker); err != nil {
			msg := fmt.Sprintf("failed to get the coprocessor worker group %s: %v", key, err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
		conditions = worker.Status.Conditions
		if worker.Spec.Template.Annotations[operationIDAnnotationKey] != req.OperationId {
			if worker.Spec.Template.Annotations == nil {
				worker.Spec.Template.Annotations = make(map[string]string)
			}
			worker.Spec.Template.Annotations[operationIDAnnotationKey] = req.OperationId
			group = &worker
		}

	default:
		err := fmt.Errorf("unsupported component type: %s", req.TargetName)
		logger.Error(err.Error())
		return &apiv1.RestartResp{Message: err.Error()}, err
	}

	if group != nil {
		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 &apiv1.RestartResp{Message: err.Error()}, err
		}

		if err := s.dataplaneCli.Update(ctx, group); err != nil {
			msg := fmt.Sprintf("failed to update the group: %v", err)
			logger.Error(msg)
			return &apiv1.RestartResp{Message: err.Error()}, err
		}
	}
	return &apiv1.RestartResp{Message: "success"}, 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{K8sComponentLabelKey},
		); 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[K8sComponentLabelKey]]; 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
				}
			}
		}

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

	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
}

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", namespacePrefix, 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 componentTypeForReplaceNodePod(&pod) == "" {
		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 := s.dataplaneCli.Update(ctx, &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
}