package server

import (
	"context"
	"math"
	"strconv"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	client "sigs.k8s.io/controller-runtime/pkg/client"
)

func (s *serverImpl) CheckPods(
	ctx context.Context,
	sharedPoolId,
	clusterId string,
	componentType entity.ComponentType,
	from string,
) (*entity.CheckPodsResponse, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "sharedPoolId is required")
	}
	if from == "" {
		return nil, status.Errorf(codes.InvalidArgument, "from is required")
	}
	fromTS, err := strconv.ParseInt(from, 10, 64)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "from is invaild: %s", err)
	}

	if componentType == apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
		return nil, status.Errorf(codes.InvalidArgument, "component type is required")
	}

	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return nil, err
	}
	if componentType == apiv1.ComponentType_COPROCESSOR_WORKER {
		if err := ensureCoprocessorWorkerComponentExists(sharedPool); err != nil {
			return nil, err
		}
	}
	remoteCli, err := s.dataplaneManager.GetClient(sharedPool.GetDataplaneName())
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get remote client for shared pool: %s", err)
	}
	namespace := "tidb" + sharedPool.GetId()
	pods := new(corev1.PodList)
	switch componentType {
	case apiv1.ComponentType_TIDB:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelUserClusterKey: clusterId,
				schemav1.LabelComponentKey:   "tidb",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tidb pods: %s", err)
		}
	case apiv1.ComponentType_TIDB_WORKER:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey:    "tidb",
				"app.kubernetes.io/component": "worker-tidb",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tidb-worker pods: %s", err)
		}
	case apiv1.ComponentType_TIKV:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: "tikv",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tikv pods: %s", err)
		}
	case apiv1.ComponentType_TIKV_WORKER:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey:    "tikv-worker",
				"app.kubernetes.io/component": "tikv-worker",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tikv-worker pods: %s", err)
		}
	case apiv1.ComponentType_COPROCESSOR_WORKER:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey:    "tikv-worker",
				"app.kubernetes.io/component": "coprocessor-worker",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list coprocessor-worker pods: %s", err)
		}
	case apiv1.ComponentType_PD:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: "pd",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list pd pods: %s", err)
		}
	case apiv1.ComponentType_PD_TSO:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: "tso",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list pd tso pods: %s", err)
		}
	case apiv1.ComponentType_PD_SCHEDULING:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: "scheduling",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list pd scheduling pods: %s", err)
		}
	case apiv1.ComponentType_TIPROXY:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelUserClusterKey: clusterId,
				schemav1.LabelComponentKey:   "tiproxy",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tiproxy pods: %s", err)
		}
	case apiv1.ComponentType_TIFLASH_COMPUTE:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey:    "tiflash",
				"app.kubernetes.io/component": "compute-tiflash",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tiflash compute pods: %s", err)
		}
	case apiv1.ComponentType_TIFLASH_WRITE:
		if err := remoteCli.List(ctx, pods, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey:    "tiflash",
				"app.kubernetes.io/component": "write-tiflash",
			}),
		}); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list tiflash write pods: %s", err)
		}
	default:
		return nil, status.Errorf(codes.InvalidArgument, "component %s is not supported", componentType)
	}
	out := new(entity.CheckPodsResponse)
	out.Recreate = checkRecreated(pods, fromTS)
	out.ZoneBalance = checkZoneBalance(pods)
	return out, nil
}

func checkRecreated(pods *corev1.PodList, fromTS int64) bool {
	for i := range pods.Items {
		pod := pods.Items[i]
		if pod.CreationTimestamp.Unix() >= fromTS {
			return true
		}
	}
	return false
}

func checkZoneBalance(pods *corev1.PodList) bool {
	zoneCnt := make(map[string]int)
	for i := range pods.Items {
		pod := pods.Items[i]
		if zone, ok := pod.Spec.NodeSelector["topology.kubernetes.io/zone"]; ok {
			zoneCnt[zone] += 1
		}
	}
	if len(zoneCnt) == 0 {
		return true
	}
	maxCnt := 0
	minCnt := math.MaxInt
	for _, cnt := range zoneCnt {
		maxCnt = max(cnt, maxCnt)
		minCnt = min(cnt, minCnt)
	}
	return maxCnt-minCnt <= 1
}