package extratidb

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"regexp"
	"sort"
	"strconv"

	oplabel "github.com/pingcap/tidb-operator/pkg/apis/label"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

// GetNodeCRLabelSelector returns the label selector for the extra TiDB related CRs.
func GetNodeCRLabelSelector(clusterID string) labels.Selector {
	return labels.SelectorFromSet(map[string]string{
		"tidbcloud-cluster": clusterID,
		"tidb-component":    "extra-tidb",
	})
}

// GetNodeCRLabels returns the labels for the extra TiDB related CRs.
func GetNodeCRLabels(clusterID, zone string) map[string]string {
	return map[string]string{
		"tidbcloud-cluster": clusterID,
		// Do not use `tidb`, so it will NOT be selected by other operations.
		"tidb-component":    "extra-tidb",
		consts.ZoneLabelKey: zone,
	}
}

// GetExtraTiDBZone returns the zone with the least TiDB nodes.
func GetExtraTiDBZone(ctx context.Context, cli client.Client, clusterID string, zones []string) (extraTiDBZone string, err error) {
	// Get all TiDB nodes.
	var tidbNodeList corev1.NodeList
	if err = cli.List(ctx, &tidbNodeList,
		client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				// The extra tidb node does not have these labels.
				"tidbcloud-cluster": clusterID,
				"tidb-component":    "tidb",
			}),
		},
	); err != nil || len(tidbNodeList.Items) == 0 {
		return "", fmt.Errorf("failed to list cluster %s's tidb nodes: %w", clusterID, err)
	}

	countPerZones := map[string]int{}
	for _, node := range tidbNodeList.Items {
		if zone, ok := node.Labels[consts.ZoneLabelKey]; ok {
			countPerZones[zone]++
		}
	}
	sort.Strings(zones)

	leastZoneCount := math.MaxInt
	for _, zone := range zones {
		if countPerZones[zone] < leastZoneCount {
			leastZoneCount = countPerZones[zone]
			extraTiDBZone = zone
		}
	}

	return extraTiDBZone, nil
}

func extraTiDBNodeGroupName(clusterID, zone string) string {
	return fmt.Sprintf("extra-tidb-%s-%s", clusterID, zone)
}

// MutateTCToScaleInExtraTiDBPod scale in the extra tidb pod.
// It first gets the pod ordinal from the pod name, then appends the pod ordinal to the delete slots.
func MutateTCToScaleInExtraTiDBPod(tc *opv1alpha1.TidbCluster, podName string, tidbReplicas int32) error {
	// Extract the pod ordinal from the pod name
	podOrdinal, err := extractNumber(podName)
	if err != nil {
		return fmt.Errorf("failed to extract pod ordinal from pod name %s: %w", podName, err)
	}

	// Get current delete slots value from tc annotations
	deleteSlots := util.GetAnnotation(tc, oplabel.AnnTiDBDeleteSlots)

	// Try to append the extracted number to the current delete slots
	newSlots, err := appendNumberToStringList(deleteSlots, podOrdinal)
	if err != nil {
		return fmt.Errorf("failed to append pod ordinal %d to delete slots %s: %w", podOrdinal, deleteSlots, err)
	}

	// Set new delete slots value to tc annotations
	util.SetAnnotation(tc, oplabel.AnnTiDBDeleteSlots, newSlots)

	// Adjust the TidbCluster tc replicas number
	tc.Spec.TiDB.Replicas = tidbReplicas

	return nil
}

func extractNumber(s string) (int32, error) {
	re := regexp.MustCompile("[0-9]+")
	result := re.FindString(s)
	parsedResult, err := strconv.ParseInt(result, 10, 32)
	return int32(parsedResult), err
}

func appendNumberToStringList(numsList string, numToAdd int32) (string, error) {
	// Check if the input string is empty
	if numsList == "" {
		// If the input string is empty, create a new slice with the number to add
		numsSlice := []int32{numToAdd}
		updatedNums, err := json.Marshal(numsSlice)
		if err != nil {
			return "", err
		}
		return string(updatedNums), nil
	}

	// Unmarshal the input string to a slice of integers
	var numsSlice []int32
	err := json.Unmarshal([]byte(numsList), &numsSlice)
	if err != nil {
		return "", err
	}

	// Check if the slice already contains the number
	for _, num := range numsSlice {
		if num == numToAdd {
			return numsList, nil
		}
	}

	// If the slice does not contain the number, append it to the slice
	numsSlice = append(numsSlice, numToAdd)
	updatedNums, err := json.Marshal(numsSlice)
	if err != nil {
		return "", err
	}
	return string(updatedNums), nil
}