package commontasks

import (
	"context"
	"fmt"

	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"

	clusterserverv1 "github.com/tidbcloud/observability/agent-control-server/api/clusterserver/v1"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/metrics"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

const (
	defaultMinSpec = "default"
)

// ReconcileAutoResourcing reconciles the desired spec annotation based on the node count.
// It sets the desired spec annotation on the agent object if it differs from the current value.
// It also sets the previous spec and spec scale status annotations when a change occurs.
func ReconcileAutoResourcing(
	ctx context.Context,
	specs config.CustomerAgentSpecs,
	agent client.Object,
	clusterName string,
	agentName string,
	defaultMaxSpec string,
	matchingLabels map[string]string,
	cli client.Client,
) error {
	util.SetAnnotationIfEmpty(agent, consts.AnnotationKeyO11yDesiredSpec, defaultMinSpec)
	minSpec := util.SetAnnotationIfEmpty(agent, consts.AnnotationKeyO11yMinSpec, defaultMinSpec)
	maxSpec := util.SetAnnotationIfEmpty(agent, consts.AnnotationKeyO11yMaxSpec, defaultMaxSpec)
	logger := ctrl.Log.WithValues("autoscaler", "horizontal")

	nodeCount, err := countComponentReplicas(ctx, cli, matchingLabels)
	if err != nil {
		return err
	}

	desiredSpec := determineSpec(specs, defaultMaxSpec, minSpec, maxSpec, nodeCount)
	prevSpec := util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec)
	if prevSpec != desiredSpec {
		metrics.IncScaledResourceCounter([]string{agentName, clusterName, compareSpec(specs, prevSpec, desiredSpec)})
		logger.Info(fmt.Sprintf("agent (%s) scales spec from (%s) to (%s) for cluster (%s), nodeCount=%d", agent, prevSpec, desiredSpec, clusterName, nodeCount))
		util.SetAnnotation(agent, consts.AnnotationKeyO11yPreviousSpec, util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec))
		util.SetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec, desiredSpec)
		util.SetAnnotation(agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusPendingScale)
	}
	return nil
}

func countComponentReplicas(ctx context.Context, cli client.Client, matchingLabels map[string]string) (uint32, error) {
	var clusters clusterserverv1.ClusterList
	if err := cli.List(ctx, &clusters, client.MatchingLabels(matchingLabels)); err != nil {
		return 0, fmt.Errorf("failed to list clusters: %w", err)
	}
	if len(clusters.Items) == 0 {
		return 0, nil
	}

	var count uint32
	for _, cluster := range clusters.Items {
		spec := cluster.Spec.Cluster
		count += uint32(spec.PD.Replicas) + uint32(spec.TiDB.Replicas) + uint32(spec.TiKV.Replicas) + uint32(spec.TiFlash.Replicas)
	}
	return count, nil
}

// compareSpec compares two spec names and determines the scaling action.
func compareSpec(specs config.CustomerAgentSpecs, prev, next string) string {
	if specs.Index(prev) > specs.Index(next) {
		return consts.ScaleActionDown
	}
	if specs.Index(prev) < specs.Index(next) {
		return consts.ScaleActionUp
	}
	return ""
}

// determineSpec determines the desired spec based on the node count and spec range.
// It returns the desired spec name.
func determineSpec(specs config.CustomerAgentSpecs, defaultMaxSpec, minSpec, maxSpec string, nodeCount uint32) string {
	if minSpec == "" {
		minSpec = defaultMinSpec
	}
	if maxSpec == "" {
		maxSpec = defaultMaxSpec
	}
	desiredSpec := ""
	arrivedMin := false
	arrivedMax := false
	for _, spec := range specs {
		if spec.Name == minSpec {
			arrivedMin = true
		}
		if spec.Name == maxSpec {
			arrivedMax = true
		}
		if !arrivedMin {
			continue
		}
		if nodeCount <= spec.MaxNodes {
			desiredSpec = spec.Name
			break
		}
		if arrivedMax {
			break
		}
	}
	if desiredSpec == "" {
		desiredSpec = minSpec
	}
	return desiredSpec
}