package cluster

import (
	"encoding/json"
	"strconv"
	"strings"

	"github.com/pingcap/errors"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/mixeddeploypd"
	"github.com/tidbcloud/infra-api/pkg/k8s"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
)

func (r *reconcileContext) scaleCluster(tc *v1alpha1.TidbCluster) error {
	desired := r.Target.Spec.Cluster
	zones, err := r.Provider.FindRegionZones(r.Network.Spec.Region)
	if err != nil {
		return errors.Annotate(err, "FindRegionZones")
	}
	if len(zones) > 1 && desired.TiKV.Replicas < int(tc.Spec.TiKV.Replicas) {
		// make sure the tikv instances are distributed evenly across multiple zones
		tikvTopo, err := k8s.GetComponentTopology(r.ctx, r.shootClient, types.NamespacedName{Namespace: tc.Namespace, Name: tc.Name}, cloudv1.ComponentTiKV)
		if err != nil {
			return errors.Annotate(err, "GetComponentTopology")
		}
		if len(tikvTopo.UnScheduled) > 0 {
			// this is an exception since we only scale cluster when cluster is normal
			return errors.New("unscheduled TiKV pods found, stop scale-in")
		}
		if tikvTopo.Replicas < int(tc.Spec.TiKV.Replicas) {
			return errors.Errorf("TiKV scale-out did not complete, tc spec %d, tikv pods %d, stop scale-in", int(tc.Spec.TiKV.Replicas), tikvTopo.Replicas)
		}
		scaleInCount := int(tc.Spec.TiKV.Replicas) - desired.TiKV.Replicas
		tikvTopoCopy := k8s.PodsByZone{}
		for k, v := range tikvTopo.Scheduled {
			tikvTopoCopy[k] = append([]string{}, v...)
		}
		// find victims
		victimOrdinals := sets.NewInt()
		for i := 0; i < scaleInCount; i++ {
			victim, err := scaleInOnePod(zones, tikvTopoCopy)
			if err != nil {
				return errors.Annotate(err, "scaleInOnePod")
			}
			victimOrdinals.Insert(victim.ordinal)
		}
		currentSlots := getTiKVDeleteSlots(tc)
		intersection := victimOrdinals.Intersection(currentSlots)
		if intersection.Len() > 0 {
			return errors.Errorf("victims has intersection with current delete-slots %+v, which means the current observation is stale", intersection.List())
		}
		newSlots := currentSlots.Union(victimOrdinals)
		err = setTiKVDeleteSlots(tc, newSlots)
		if err != nil {
			return errors.Annotate(err, "setTiKVDeleteSlots")
		}
	}
	// TODO(aylei): investigate whether tiflash requires this

	// Here is for:
	// * Changing from mixed-deploy (PD with TiKV) to normal-deployment, if needed
	//
	// * We do this checking no matter:
	//   - Mixed-deploy is enabled or not
	//   - The cluster is affected by mixed-deploy or not
	//
	// * If we only do the checking for affected clusters (and when mixed-deploy enabled),
	//   When a cluster is affected then become not affected by changing infra-api config file, it might:
	//   - Stay on mixed-deploy mode forever if no scaling happens, which not match the expectation of configs
	//   - PD pods might have complicated status, in Scaling and Upgrading at the same time:
	//       * Scaling because `node-selector` is changed, PD pods need to move out from TiKV nodes
	//       * Upgrading because resource spec is changed, PD pods need to be rescheduled, which will cause restarting (currently)
	//   - Solve these complications by doing the restore-checking in no matter what situations
	//       * In mixed-deploy implementation, leaving mixed-deploy is splitted to multiply steps
	tcName := types.NamespacedName{
		Namespace: tc.Namespace,
		Name:      tc.Name,
	}.String()
	expectIsMixedDeploy, err := mixeddeploypd.ShouldMixedDeployPDWithTiKV(&desired, "scaling", r.Target.Namespace, tcName,
		r.Provider, tc.ObjectMeta.CreationTimestamp.Time, r.PDMixedDeployConfig)
	if err != nil {
		return err
	}
	if expectIsMixedDeploy {
		err = mixeddeploypd.ChangeTCForMixedDeploy(tc, r.Target, r.ConfigProvider, r.Provider, r.GetComponentNodeLabels)
	} else {
		err = mixeddeploypd.RestoreTCForNonMixedDeploy(tc, r.Target, r.ConfigProvider, r.Provider, r.shootClient, r.GetComponentNodeLabels)
	}
	if err != nil {
		return err
	}

	// set all components to desired replicas
	return r.overlayReplicas(tc)
}

type statefulPod struct {
	name    string
	ordinal int
}

func getTiKVDeleteSlots(tc *v1alpha1.TidbCluster) sets.Int {
	deleteSlots := sets.NewInt()
	annotations := tc.GetAnnotations()
	if annotations == nil {
		return deleteSlots
	}
	value, ok := annotations[label.AnnTiKVDeleteSlots]
	if !ok {
		return deleteSlots
	}
	var slice []int
	err := json.Unmarshal([]byte(value), &slice)
	if err != nil {
		return deleteSlots
	}
	deleteSlots.Insert(slice...)
	return deleteSlots
}

func setTiKVDeleteSlots(tc *v1alpha1.TidbCluster, slots sets.Int) error {
	annotations := tc.GetAnnotations()
	if slots == nil || slots.Len() == 0 {
		// clear
		delete(annotations, label.AnnTiDBDeleteSlots)
	} else {
		// set
		b, err := json.Marshal(slots.List())
		if err != nil {
			return err
		}
		if annotations == nil {
			annotations = make(map[string]string)
		}
		annotations[label.AnnTiKVDeleteSlots] = string(b)
	}
	tc.SetAnnotations(annotations)
	return nil
}

// scaleInOnePod select the victim to be scaled-in and remove it from current topology
func scaleInOnePod(zones []string, spreads k8s.PodsByZone) (*statefulPod, error) {
	maxZoneIndex := 0
	maxZoneCount := 0
	// find the max zone to scale-in, latter zone wins on tie
	for i := range zones {
		if len(spreads[zones[i]]) >= maxZoneCount {
			maxZoneCount = len(spreads[zones[i]])
			maxZoneIndex = i
		}
	}
	maxZone := zones[maxZoneIndex]
	maxZonePods := spreads[maxZone]
	victim := &statefulPod{ordinal: -1}
	victimIndex := 0
	// find the max ordinal to scale-in in the max zone
	for i, podName := range maxZonePods {
		pod, err := parseStatefulPod(podName)
		if err != nil {
			return nil, errors.Annotatef(err, "%s is not a statefulset pod", podName)
		}
		if pod.ordinal > victim.ordinal {
			victim = pod
			victimIndex = i
		}
	}
	if victim.ordinal < 0 {
		return nil, errors.New("no pod to scale-in")
	}
	// rm victim from current topology
	spreads[maxZone] = append(maxZonePods[:victimIndex], maxZonePods[victimIndex+1:]...)
	return victim, nil
}

func parseStatefulPod(name string) (*statefulPod, error) {
	splits := strings.Split(name, "-")
	// last part should be ordinal for statefulset pod
	ordinalStr := splits[len(splits)-1]
	ordinal, err := strconv.Atoi(ordinalStr)
	if err != nil {
		return nil, err
	}
	return &statefulPod{
		name:    name,
		ordinal: ordinal,
	}, nil
}