package podautoscaler

import (
	"context"
	"errors"
	"fmt"
	"math"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/sets"
	resourcehelpers "k8s.io/component-helpers/resource"
	resourceclient "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1beta1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
)

const (
	// defaultMemoryInitializationPeriod is the period after pod start when CPU samples
	// might be skipped.
	defaultCpuInitializationPeriod = time.Minute * 5

	// defaultDelayOfInitialReadinessStatus is period after pod start during which readiness
	// changes are treated as readiness being set for the first time. The only effect of this is that
	// HPA will disregard CPU samples from unready pods that had last readiness change during that
	// period.
	defaultDelayOfInitialReadinessStatus = time.Second * 30

	// defaultResultQueueSize is the size of the queue to store calculation results
	defaultResultQueueSize = 15
)

// Tolerances is the tolerance on the ratio between the current and desired
// metric value under which no updates are made to the desired number of
// replicas (e.g. 0.01 for 1%). Must be greater than or equal to zero. If not
// set, the default cluster-wide tolerance is applied (by default 10%).
//
// For example, if autoscaling is configured with a memory consumption target of 100Mi,
// and scale-down and scale-up tolerances of 5% and 1% respectively, scaling will be
// triggered when the actual consumption falls below 95Mi or exceeds 101Mi.
type Tolerances struct {
	ScaleDown float64
	ScaleUp   float64
}

// DefaultTolerances returns the default tolerances for scaling up and down
func DefaultTolerances() Tolerances {
	return Tolerances{
		ScaleDown: 0.05,
		ScaleUp:   0.05,
	}
}

func (t Tolerances) String() string {
	return fmt.Sprintf("[down:%.1f%%, up:%.1f%%]", t.ScaleDown*100., t.ScaleUp*100.)
}

// IsWithin returns true if the given usage ratio is within the tolerances.
// That is, if the usage ratio is between (1 - ScaleDown) and (1 + ScaleUp).
func (t Tolerances) IsWithin(usageRatio float64) bool {
	return (1.0-t.ScaleDown) <= usageRatio && usageRatio <= (1.0+t.ScaleUp)
}

// resourceCalculator bundles all needed information to calculate the target amount of replicas
type resourceCalculator struct {
	metricsClient    *metricsClient
	clientReader     client.Reader
	tolerances       Tolerances
	clusterName      string
	namespace        string
	selector         labels.Selector
	memq             *queueManager
	cpuq             *queueManager
	autoscalerConfig config.AutoScaler
}

// newResourceCalculator creates a new ResourceCalculator and passes all necessary information to the new instance
func newResourceCalculator(
	metricsInterface resourceclient.MetricsV1beta1Interface,
	clientReader client.Reader,
	clusterName string,
	namespace string,
	matchingLabels map[string]string,
	autoscalerConfig config.AutoScaler,
) *resourceCalculator {
	return &resourceCalculator{
		metricsClient:    newMetricsClient(metricsInterface),
		clientReader:     clientReader,
		clusterName:      clusterName,
		namespace:        namespace,
		selector:         labels.SelectorFromSet(matchingLabels),
		memq:             newQueueManager(defaultResultQueueSize),
		cpuq:             newQueueManager(defaultResultQueueSize),
		tolerances:       DefaultTolerances(),
		autoscalerConfig: autoscalerConfig,
	}
}

func (rc *resourceCalculator) String() string {
	return fmt.Sprintf("namespace=%s,%s", rc.namespace, rc.selector.String())
}

// ScaleResult contains the result of a scaling calculation for a specific resource
type ScaleResult struct {
	// GroupKey is the key of the object being scaled
	GroupKey GroupKey
	// Resource is the resource being scaled (cpu/memory)
	Resource corev1.ResourceName
	// ActionResult is the result of the scaling action calculation
	ActionResult ActionResult
	// Pods is the number of pods considered in the calculation
	Pods int
}

// CalculateScaleResult calculates the scale result for both memory and CPU resources
// It refreshes the raw results for both resources and then computes the scale results
// using the provided auto scale rule. The results are returned as a slice of ScaleResult.
// If any error occurs during the process, it is returned.
func (rc *resourceCalculator) CalculateScaleResult(ctx context.Context) ([]ScaleResult, error) {
	var err error
	err = rc.refreshResourceRawResult(ctx, rc.autoscalerConfig.TargetUtilization, corev1.ResourceCPU)
	if err != nil {
		if errors.Is(err, ErrNoMetricsFound) {
			logger.Info(fmt.Sprintf("cluster (%s) no cpu metrics found for selector (%s)", rc.clusterName, rc.String()))
			return nil, nil
		}
		return nil, err
	}
	err = rc.refreshResourceRawResult(ctx, rc.autoscalerConfig.TargetUtilization, corev1.ResourceMemory)
	if err != nil {
		if errors.Is(err, ErrNoMetricsFound) {
			logger.Info(fmt.Sprintf("cluster (%s) no memory metrics found for selector (%s)", rc.clusterName, rc.String()))
			return nil, nil
		}
		return nil, err
	}

	return rc.doCalculateScaleResult(), nil
}

func (rc *resourceCalculator) doCalculateScaleResult() []ScaleResult {
	// Calculate scale results for CPU and memory resources
	merged := make(map[GroupKey][]ScaleResult)
	for _, sr := range calcScaleResult(rc.autoscalerConfig, rc.cpuq, corev1.ResourceCPU) {
		merged[sr.GroupKey] = append(merged[sr.GroupKey], sr)
	}
	for _, sr := range calcScaleResult(rc.autoscalerConfig, rc.memq, corev1.ResourceMemory) {
		merged[sr.GroupKey] = append(merged[sr.GroupKey], sr)
	}

	// Merge scale results, prioritizing scale-up actions
	// If there are multiple scale-down actions, only one is kept
	// to avoid excessive scaling down
	var ans []ScaleResult
	for _, srs := range merged {
		var down int
		var added bool
		// Scaling up only needs to meet either CPU or memory conditions,
		// while scaling down requires both CPU and memory conditions to be met.
		// (Scaling up is proactive, scaling down is conservative)
		for _, sr := range srs {
			if sr.ActionResult.Action == consts.ScaleActionUp {
				ans = append(ans, sr)
				added = true
				break
			}
			if sr.ActionResult.Action == consts.ScaleActionDown {
				down++
			}
		}
		// Only scale down if BOTH CPU and memory indicate down (conservative approach)
		if !added && down > 1 {
			ans = append(ans, srs[0])
		}
	}

	for k := range merged {
		rc.cpuq.Reset(k)
		rc.memq.Reset(k)
	}
	return ans
}

// ActionResult contains the result of a scaling action calculation for a specific resource
type ActionResult struct {
	// Action is the type of scaling action: "up" or "down"
	Action string
	// Samples is the number of samples considered for the scaling decision
	Samples int
	// Hits is the number of samples that met the scaling criteria
	Hits int32
	// TargetUtilization is the target utilization percentage for the resource
	TargetUtilization int32
	// AverageUtilization is the average utilization percentage across the samples that met the scaling criteria
	AverageUtilization int32
}

// refreshResourceRawResult refreshes the raw results for the specified resource
// It retrieves the resource metrics, lists the pods, and calculates the replica count,
// utilization, and raw average value for each pod group. The results are stored in the
// appropriate queue manager (CPU or memory) based on the resource type.
func (rc *resourceCalculator) refreshResourceRawResult(ctx context.Context, targetUtilization int32, resource corev1.ResourceName) error {
	metrics, ts, err := rc.metricsClient.GetResourceMetric(ctx, resource, rc.namespace, rc.selector)
	if err != nil {
		return err
	}

	podList, err := rc.listPods(ctx)
	if err != nil {
		return fmt.Errorf("unable to get pods while calculating replica count: %v", err)
	}
	if len(podList) == 0 {
		return fmt.Errorf("no pods returned by selector while calculating replica count")
	}

	podGroups, metricGroups := groupByPodMetrics(podList, metrics)
	for groupKey, pods := range podGroups {
		replicaCount, utilization, rawAverageValue, err := getResourceReplicasGroup(rc.tolerances, pods, metricGroups[groupKey], int32(len(pods)), targetUtilization, resource)
		if err != nil {
			return err
		}

		q := rc.memq
		if resource == corev1.ResourceCPU {
			q = rc.cpuq
		}
		q.Put(groupKey, RawResult{
			Resource:        resource,
			Time:            ts,
			ReplicaCount:    replicaCount,
			Utilization:     utilization,
			RawAverageValue: rawAverageValue,
			Pods:            len(pods),
		})
	}
	return nil
}

// listPods lists the pods in the specified namespace that match the selector.
func (rc *resourceCalculator) listPods(ctx context.Context) ([]*corev1.Pod, error) {
	var pods corev1.PodList
	if err := rc.clientReader.List(ctx, &pods, client.InNamespace(rc.namespace), client.MatchingLabelsSelector{Selector: rc.selector}); err != nil {
		return nil, fmt.Errorf("unable to list pods in namespace %s: %v", rc.namespace, err)
	}

	var podList []*corev1.Pod
	for i := 0; i < len(pods.Items); i++ {
		podList = append(podList, &pods.Items[i])
	}
	return podList, nil
}

// calcScaleResult calculates the scale result for the given resource using the provided queue manager.
// It iterates over the object keys in the queue manager, calculates the action result for each key,
func calcScaleResult(autoscalerConfig config.AutoScaler, qm *queueManager, resource corev1.ResourceName) []ScaleResult {
	var srs []ScaleResult //nolint:prealloc
	for _, groupKey := range qm.GroupKeys() {
		rawRets := qm.List(groupKey)
		if len(rawRets) == 0 {
			continue
		}
		ar := calcActionResult(autoscalerConfig, rawRets)
		if ar == nil {
			continue
		}

		// check if the last reset timestamp is within the frozen period
		lastRst := qm.LastResetTimestamp(groupKey)
		if lastRst > 0 && time.Now().Unix()-lastRst < autoscalerConfig.FrozenSeconds {
			continue
		}

		srs = append(srs, ScaleResult{
			GroupKey:     groupKey,
			Resource:     resource,
			ActionResult: *ar,
			Pods:         rawRets[len(rawRets)-1].Pods, // use the latest pod count
		})
	}
	return srs
}

// getResourceReplicasGroup calculates the desired number of replicas for a group of pods
// based on the provided tolerances, pod metrics, current replicas, target utilization, and resource type.
// It returns the calculated replica count, current utilization, raw average value, and any error encountered.
// The function considers ready, unready, missing, and ignored pods in its calculations.
func getResourceReplicasGroup(tolerances Tolerances, pods []*corev1.Pod, metrics PodMetricsInfo, currentReplicas int32, targetUtilization int32, resource corev1.ResourceName) (replicaCount int32, utilization int32, rawAverageValue int64, err error) {
	readyPodCount, unreadyPods, missingPods, ignoredPods := groupPods(pods, metrics, resource, defaultCpuInitializationPeriod, defaultDelayOfInitialReadinessStatus)
	removeMetricsForPods(metrics, ignoredPods)
	removeMetricsForPods(metrics, unreadyPods)
	if len(metrics) == 0 {
		return 0, 0, 0, fmt.Errorf("did not receive metrics for targeted pods (pods might be unready)")
	}

	requests, err := calculateRequests(pods, resource)
	if err != nil {
		return 0, 0, 0, err
	}

	usageRatio, utilization, rawAverageValue, err := getResourceUtilizationRatio(metrics, requests, targetUtilization)
	if err != nil {
		return 0, 0, 0, err
	}

	scaleUpWithUnready := len(unreadyPods) > 0 && usageRatio > 1.0
	if !scaleUpWithUnready && len(missingPods) == 0 {
		if tolerances.IsWithin(usageRatio) {
			// return the current replicas if the change would be too small
			return currentReplicas, utilization, rawAverageValue, nil
		}

		// if we don't have any unready or missing pods, we can calculate the new replica count now
		return int32(math.Ceil(usageRatio * float64(readyPodCount))), utilization, rawAverageValue, nil
	}

	if len(missingPods) > 0 {
		if usageRatio < 1.0 {
			// on a scale-down, treat missing pods as using 100% (all) of the resource request
			// or the utilization target for targets higher than 100%
			fallbackUtilization := int64(max(100, targetUtilization))
			for pod := range missingPods {
				metrics[pod] = PodMetric{Value: requests[pod] * fallbackUtilization / 100}
			}
		} else if usageRatio > 1.0 {
			// on a scale-up, treat missing pods as using 0% of the resource request
			for pod := range missingPods {
				metrics[pod] = PodMetric{Value: 0}
			}
		}
	}

	if scaleUpWithUnready {
		// on a scale-up, treat unready pods as using 0% of the resource request
		for pod := range unreadyPods {
			metrics[pod] = PodMetric{Value: 0}
		}
	}

	// re-run the utilization calculation with our new numbers
	newUsageRatio, _, _, err := getResourceUtilizationRatio(metrics, requests, targetUtilization)
	if err != nil {
		return 0, utilization, rawAverageValue, err
	}

	if tolerances.IsWithin(newUsageRatio) || (usageRatio < 1.0 && newUsageRatio > 1.0) || (usageRatio > 1.0 && newUsageRatio < 1.0) {
		// return the current replicas if the change would be too small,
		// or if the new usage ratio would cause a change in scale direction
		return currentReplicas, utilization, rawAverageValue, nil
	}

	newReplicas := int32(math.Ceil(newUsageRatio * float64(len(metrics))))
	if (newUsageRatio < 1.0 && newReplicas > currentReplicas) || (newUsageRatio > 1.0 && newReplicas < currentReplicas) {
		// return the current replicas if the change of metrics length would cause a change in scale direction
		return currentReplicas, utilization, rawAverageValue, nil
	}

	// return the result, where the number of replicas considered is
	// however many replicas factored into our calculation
	return newReplicas, utilization, rawAverageValue, nil
}

// calcActionResult calculates the action result based on the scale rule and raw results.
// It returns nil if the number of raw results is less than the scale rule samples.
// Otherwise, it calculates the action result based on the scale rule and raw results.
func calcActionResult(autoscalerConfig config.AutoScaler, rets []RawResult) *ActionResult {
	l := len(rets)
	if l < autoscalerConfig.Samples {
		return nil
	}

	var up, down int32
	var upSum, downSum int32
	for i := l - 1; i >= l-autoscalerConfig.Samples; i-- {
		utilization := rets[i].Utilization
		if utilization <= 0 {
			continue
		}
		if utilization >= autoscalerConfig.TargetUtilization {
			up++
			upSum += utilization
		}
		// Down action is only triggered when the utilization is below 1/4 of the target utilization.
		// This is to avoid scaling down when the utilization is still high.
		//
		// Made scale-down more conservative to prevent premature resource reduction
		if utilization < autoscalerConfig.TargetUtilization/4 {
			down++
			downSum += utilization
		}
	}

	// Prioritize scaling up
	if up >= autoscalerConfig.Hits {
		return &ActionResult{
			Action:             consts.ScaleActionUp,
			Hits:               up,
			Samples:            autoscalerConfig.Samples,
			TargetUtilization:  autoscalerConfig.TargetUtilization,
			AverageUtilization: upSum / up,
		}
	}
	if down >= autoscalerConfig.Hits {
		return &ActionResult{
			Action:             consts.ScaleActionDown,
			Hits:               down,
			Samples:            autoscalerConfig.Samples,
			TargetUtilization:  autoscalerConfig.TargetUtilization,
			AverageUtilization: downSum / down,
		}
	}
	return nil
}

type GroupKey struct {
	Namespace    string
	RefNamespace string
	RefName      string
}

// groupByPodMetrics groups pods and their metrics by the reference namespace and name
// extracted from the pod labels.
// only pods with both ref namespace and ref name labels are considered.
func groupByPodMetrics(pods []*corev1.Pod, podMetrics PodMetricsInfo) (map[GroupKey][]*corev1.Pod, map[GroupKey]PodMetricsInfo) {
	selectPods := make(map[client.ObjectKey]*corev1.Pod)
	for i := 0; i < len(pods); i++ {
		pod := pods[i]
		selectPods[client.ObjectKey{Namespace: pod.Namespace, Name: pod.Name}] = pod
	}

	// Group pod metrics and pods by their reference namespace and name
	podMetricsGroup := make(map[GroupKey]PodMetricsInfo)
	podsGroup := make(map[GroupKey][]*corev1.Pod)
	for objectKey, podMetric := range podMetrics {
		pod := selectPods[objectKey]
		if pod == nil {
			continue
		}

		refNamespace := pod.Labels[consts.AgentKeyRefNamespace]
		refName := pod.Labels[consts.AgentKeyRefName]
		if refNamespace == "" || refName == "" {
			continue
		}

		groupKey := GroupKey{
			Namespace:    objectKey.Namespace,
			RefNamespace: refNamespace,
			RefName:      refName,
		}
		if _, ok := podMetricsGroup[groupKey]; !ok {
			podMetricsGroup[groupKey] = make(PodMetricsInfo)
		}
		podMetricsGroup[groupKey][objectKey] = podMetric
		podsGroup[groupKey] = append(podsGroup[groupKey], pod)
	}

	return podsGroup, podMetricsGroup
}

// groupPods classifies pods into ready, unready, missing metrics, and ignored pods.
func groupPods(pods []*corev1.Pod, metrics PodMetricsInfo, resource corev1.ResourceName, cpuInitializationPeriod, delayOfInitialReadinessStatus time.Duration) (readyPodCount int, unreadyPods, missingPods, ignoredPods sets.Set[client.ObjectKey]) {
	missingPods = sets.New[client.ObjectKey]()
	unreadyPods = sets.New[client.ObjectKey]()
	ignoredPods = sets.New[client.ObjectKey]()
	for _, pod := range pods {
		objKey := client.ObjectKey{Namespace: pod.Namespace, Name: pod.Name}
		if pod.DeletionTimestamp != nil || pod.Status.Phase == corev1.PodFailed {
			ignoredPods.Insert(objKey)
			continue
		}
		// Pending pods are ignored.
		if pod.Status.Phase == corev1.PodPending {
			unreadyPods.Insert(objKey)
			continue
		}
		// Pods missing metrics.
		metric, found := metrics[objKey]
		if !found {
			missingPods.Insert(objKey)
			continue
		}
		// Unready pods are ignored.
		if resource == corev1.ResourceCPU {
			var unready bool
			condition := getPodCondition(&pod.Status, corev1.PodReady)
			if condition == nil || pod.Status.StartTime == nil {
				unready = true
			} else {
				// Pod still within possible initialisation period.
				if pod.Status.StartTime.Add(cpuInitializationPeriod).After(time.Now()) {
					// Ignore sample if pod is unready or one window of metric wasn't collected since last state transition.
					unready = condition.Status == corev1.ConditionFalse || metric.Timestamp.Before(condition.LastTransitionTime.Time.Add(metric.Window))
				} else {
					// Ignore metric if pod is unready and it has never been ready.
					unready = condition.Status == corev1.ConditionFalse && pod.Status.StartTime.Add(delayOfInitialReadinessStatus).After(condition.LastTransitionTime.Time)
				}
			}
			if unready {
				unreadyPods.Insert(objKey)
				continue
			}
		}
		readyPodCount++
	}
	return
}

// calculateRequests computes the request value for each pod for the specified resource.
//
// It uses pod-level requests if pod-level requests are set on the pod.
// Otherwise, it sums the requests of all containers in the pod
// (including restartable init containers).
// It returns a map of pod names to their calculated request values.
func calculateRequests(pods []*corev1.Pod, resource corev1.ResourceName) (map[client.ObjectKey]int64, error) {
	requests := make(map[client.ObjectKey]int64, len(pods))
	for _, pod := range pods {
		var request int64
		var err error
		// Determine if we should use pod-level requests: see KEP-2837
		// https://github.com/kubernetes/enhancements/blob/master/keps/sig-node/2837-pod-level-resource-spec/README.md
		if resourcehelpers.IsPodLevelRequestsSet(pod) {
			request, err = calculatePodLevelRequests(pod, resource)
		} else {
			request, err = calculatePodRequestsFromContainers(pod, resource)
		}
		if err != nil {
			return nil, err
		}
		requests[client.ObjectKey{Namespace: pod.Namespace, Name: pod.Name}] = request
	}
	return requests, nil
}

// calculatePodLevelRequests computes the requests for the specific resource at
// the pod level.
func calculatePodLevelRequests(pod *corev1.Pod, resource corev1.ResourceName) (int64, error) {
	podLevelRequests := resourcehelpers.PodRequests(pod, resourcehelpers.PodResourcesOptions{})
	podRequest, ok := podLevelRequests[resource]
	if !ok {
		return 0, fmt.Errorf("missing pod-level request for %s in Pod %s", resource, pod.Name)
	}
	return podRequest.MilliValue(), nil
}

// calculatePodRequestsFromContainers computes the requests for the specified
// resource by summing requests from all containers in the pod.
// If a container name is specified, it uses only that container.
func calculatePodRequestsFromContainers(pod *corev1.Pod, resource corev1.ResourceName) (int64, error) {
	containers := append([]corev1.Container{}, pod.Spec.Containers...)
	for _, container := range pod.Spec.InitContainers {
		if container.RestartPolicy != nil && *container.RestartPolicy == corev1.ContainerRestartPolicyAlways {
			containers = append(containers, container)
		}
	}

	request := int64(0)
	for _, cont := range containers {
		containerRequest, ok := cont.Resources.Requests[resource]
		if !ok {
			return 0, fmt.Errorf("missing request for %s in container %s of Pod %s", resource, cont.Name, pod.Name)
		}
		request += containerRequest.MilliValue()
	}
	return request, nil
}

// removeMetricsForPods removes metrics for the given set of pods from the provided metrics map.
func removeMetricsForPods(metrics PodMetricsInfo, pods sets.Set[client.ObjectKey]) {
	for _, pod := range pods.UnsortedList() {
		delete(metrics, pod)
	}
}

// getResourceUtilizationRatio takes in a set of metrics, a set of matching requests,
// and a target utilization percentage, and calculates the ratio of
// desired to actual utilization (returning that, the actual utilization, and the raw average value)
func getResourceUtilizationRatio(metrics PodMetricsInfo, requests map[client.ObjectKey]int64, targetUtilization int32) (utilizationRatio float64, currentUtilization int32, rawAverageValue int64, err error) {
	metricsTotal := int64(0)
	requestsTotal := int64(0)
	numEntries := 0

	for pod, metric := range metrics {
		request, hasRequest := requests[pod]
		if !hasRequest {
			// we check for missing requests elsewhere, so assuming missing requests == extraneous metrics
			continue
		}

		metricsTotal += metric.Value
		requestsTotal += request
		numEntries++
	}

	// if the set of requests is completely disjoint from the set of metrics,
	// then we could have an issue where the requests total is zero
	if requestsTotal == 0 {
		return 0, 0, 0, fmt.Errorf("no metrics returned matched known pods")
	}

	currentUtilization = int32((metricsTotal * 100) / requestsTotal)
	return float64(currentUtilization) / float64(targetUtilization), currentUtilization, metricsTotal / int64(numEntries), nil
}

// getPodCondition extracts the provided condition from the PodStatus and returns that.
// Returns nil if the condition is not present.
func getPodCondition(status *corev1.PodStatus, conditionType corev1.PodConditionType) *corev1.PodCondition {
	if status == nil {
		return nil
	}
	return getPodConditionFromList(status.Conditions, conditionType)
}

// getPodConditionFromList extracts the provided condition from the given list of conditions and returns that.
// Returns nil if the condition is not present.
func getPodConditionFromList(conditions []corev1.PodCondition, conditionType corev1.PodConditionType) *corev1.PodCondition {
	if conditions == nil {
		return nil
	}
	for i := range conditions {
		if conditions[i].Type == conditionType {
			return &conditions[i]
		}
	}
	return nil
}