package podautoscaler

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/pingcap/errors"
	"go.uber.org/multierr"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"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/k8sclient/k8scliv2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/loader"
	"github.com/tidbcloud/observability/agent-control-server/pkg/metrics"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

const (
	defaultSpecName = "default"
	maxSpecName     = "large"
)

var specRanges = map[string]specRange{
	consts.AgentKubernetes: {minSpec: defaultSpecName, maxSpec: maxSpecName},
	consts.AgentSharedPool: {minSpec: defaultSpecName, maxSpec: maxSpecName},
	consts.AgentTidbGroup:  {minSpec: defaultSpecName, maxSpec: maxSpecName},
}

const (
	skipReasonFrozen  = "frozen"
	skipReasonMaxSpec = "maxspec"
	skipReasonMinSpec = "minspec"
	skipReasonOther   = "skip" // For spec not found or other reasons
)

type specSelectionResult struct {
	newSpec    string
	skipReason string
}

var logger = ctrl.Log.WithValues("autoscaler", "vertical")

// PodAutoScaler is responsible for automatically scaling pod resources based on predefined rules.
// It interacts with Kubernetes clusters to monitor and adjust resources for various agents.
// It maintains client getters and resource updaters for efficient operations across multiple clusters.
type PodAutoScaler struct {
	cli  client.Client
	done chan struct{}

	mut           sync.Mutex
	clientGetters map[string]k8scliv2.ClientGetter
	updaters      map[string]*agentResourceUpdater
	activeTs      map[string]int64

	rootConfig       config.RootConfig
	autoscalerConfig config.AutoScaler
	nextgenConfig    config.CustomerNextGen
}

func New(cli client.Client, rootConfig config.RootConfig) *PodAutoScaler {
	pas := &PodAutoScaler{
		cli:              cli,
		done:             make(chan struct{}),
		clientGetters:    make(map[string]k8scliv2.ClientGetter),
		updaters:         make(map[string]*agentResourceUpdater),
		activeTs:         make(map[string]int64),
		rootConfig:       rootConfig,
		autoscalerConfig: rootConfig.AutoScaler,
		nextgenConfig:    rootConfig.O11yConfig.CustomerConfigsV2.NextGen,
	}
	go pas.gc()
	return pas
}

// Run executes the pod auto-scaling process across all Kubernetes clusters.
// It lists all clusters, retrieves or creates client getters and resource updaters,
// and performs detection and scaling for each type of agent.
func (pas *PodAutoScaler) Run(ctx context.Context) error {
	clusters, err := pas.listKubernetesClusters(ctx)
	if err != nil {
		return err
	}

	var errs error
	for _, cluster := range clusters {
		aru, err := pas.getAgentResourceUpdater(ctx, cluster)
		if err != nil {
			return err
		}

		if err := pas.detect(ctx, aru.kubernetesVmagent); err != nil {
			errs = multierr.Append(errs, errors.Wrapf(err, "detect (%s) kubernetes/vmagent failed", cluster))
		}
		if err := pas.detect(ctx, aru.sharedpoolVmagent); err != nil {
			errs = multierr.Append(errs, errors.Wrapf(err, "detect (%s) sharedpool/vmagent failed", cluster))
		}
		if err := pas.detect(ctx, aru.sharedpoolVector); err != nil {
			errs = multierr.Append(errs, errors.Wrapf(err, "detect (%s) sharedpool/vector failed", cluster))
		}
		if err := pas.detect(ctx, aru.tidbgroupVector); err != nil {
			errs = multierr.Append(errs, errors.Wrapf(err, "detect (%s) tidbgroup/vector failed", cluster))
		}
	}
	return errs
}

func (pas *PodAutoScaler) Clean() {
	close(pas.done)
}

// listKubernetesClusters retrieves the list of Kubernetes cluster names from the KubernetesAgent resources.
func (pas *PodAutoScaler) listKubernetesClusters(ctx context.Context) ([]string, error) {
	var agents agentv2.KubernetesAgentList
	if err := pas.cli.List(ctx, &agents); err != nil {
		return nil, err
	}

	pas.mut.Lock()
	defer pas.mut.Unlock()

	clusters := make([]string, 0, len(agents.Items))
	for _, agent := range agents.Items {
		clusters = append(clusters, agent.Spec.KubernetesClusterName)
		pas.activeTs[agent.Spec.KubernetesClusterName] = time.Now().Unix()
	}
	return clusters, nil
}

func (pas *PodAutoScaler) getAgentResourceUpdater(ctx context.Context, clusterName string) (*agentResourceUpdater, error) {
	pas.mut.Lock()
	defer pas.mut.Unlock()

	if _, ok := pas.clientGetters[clusterName]; !ok {
		clientGetter, err := k8scliv2.NetworkClientGetter(
			pas.cli,
			loader.RemoteScheme(),
			consts.VendorToProvider(pas.rootConfig.O11yConfig.Vendor),
			clusterName,
			&pas.rootConfig,
		)
		if err != nil {
			return nil, err
		}
		pas.clientGetters[clusterName] = clientGetter
	}

	aru, ok := pas.updaters[clusterName]
	if ok {
		return aru, nil
	}

	clientGetter := pas.clientGetters[clusterName]
	metricsCli, err := clientGetter.GetMetricsClient(ctx)
	if err != nil {
		return nil, err
	}
	cli, err := clientGetter.GetClient(ctx)
	if err != nil {
		return nil, err
	}

	aru = &agentResourceUpdater{
		kubernetesVmagent: newResourceUpdater(
			cli,
			metricsCli,
			clusterName,
			pas.systemNamespace(),
			consts.VMAgentLabels(consts.AgentKubernetes),
			pas.autoscalerConfig,
			pas.updateKubernetesAgentResources,
			specRanges[consts.AgentKubernetes],
			consts.ComponentVMAgent,
		),
		sharedpoolVmagent: newResourceUpdater(
			cli,
			metricsCli,
			clusterName,
			"",
			consts.VMAgentLabels(consts.AgentSharedPool),
			pas.autoscalerConfig,
			pas.updateSharedPoolAgentResources,
			specRanges[consts.AgentSharedPool],
			consts.ComponentVMAgent,
		),
		sharedpoolVector: newResourceUpdater(
			cli,
			metricsCli,
			clusterName,
			"",
			consts.VectorLabels(consts.AgentSharedPool),
			pas.autoscalerConfig,
			pas.updateSharedPoolAgentResources,
			specRanges[consts.AgentSharedPool],
			consts.ComponentVector,
		),
		tidbgroupVector: newResourceUpdater(
			cli,
			metricsCli,
			clusterName,
			"",
			consts.VectorLabels(consts.AgentTidbGroup),
			pas.autoscalerConfig,
			pas.updateTidbGroupAgentResources,
			specRanges[consts.AgentTidbGroup],
			consts.ComponentVector,
		),
	}
	pas.updaters[clusterName] = aru
	return aru, nil
}

func (pas *PodAutoScaler) systemNamespace() string {
	if pas.rootConfig.IsServerlessMode() {
		return consts.ServerlessObservabilityNamespace
	}
	return consts.O11ySystemNamespace
}

func (pas *PodAutoScaler) detect(ctx context.Context, rc *resourceUpdater) error {
	srs, err := rc.CalculateScaleResult(ctx)
	if err != nil {
		return err
	}

	for _, sr := range srs {
		err = multierr.Append(err, rc.UpdateRefResources(ctx, sr))
	}
	return err
}

// gc performs garbage collection of inactive cluster client getters and updaters.
func (pas *PodAutoScaler) gc() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	f := func() {
		pas.mut.Lock()
		now := time.Now().Unix()
		for clusterName, ts := range pas.activeTs {
			// 30min
			if now-ts > 1800 {
				delete(pas.clientGetters, clusterName)
				delete(pas.updaters, clusterName)
				delete(pas.activeTs, clusterName)
				logger.Info(fmt.Sprintf("pod autoscaler gc cluster (%s) clientGetter and updater", clusterName))
			}
		}
		pas.mut.Unlock()
	}

	for {
		select {
		case <-ticker.C:
			f()
		case <-pas.done:
			logger.Info("pod autoscaler is shutting down")
			return
		}
	}
}

type specRange struct {
	minSpec string
	maxSpec string
}

type updateResourceFunc func(ctx context.Context, clusterName, component string, scaleResult ScaleResult, specRange specRange) error

func (pas *PodAutoScaler) updateKubernetesAgentResources(ctx context.Context, clusterName, component string, scaleResult ScaleResult, specRange specRange) error {
	groupKey := scaleResult.GroupKey
	var obj agentv2.KubernetesAgent
	if err := pas.cli.Get(ctx, client.ObjectKey{Namespace: groupKey.RefNamespace, Name: groupKey.RefName}, &obj); err != nil {
		return err
	}

	specs := pas.nextgenConfig.K8s.Specs
	agentComponent := consts.AgentKubernetes + "/" + component
	currSpec, nextSpec, changed, skipReason := decideSpecByScaleResult(&obj, scaleResult, specs, pas.autoscalerConfig.FrozenSeconds, specRange, component)
	if !changed {
		metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, skipReason})
		logger.Info(fmt.Sprintf("skip updating %s (%+v) spec (%s) for cluster (%s), reason=%s, scaleResult=%+v", agentComponent, groupKey, currSpec, clusterName, skipReason, scaleResult))
		return nil
	}

	metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, scaleResult.ActionResult.Action})
	logger.Info(fmt.Sprintf("%s (%+v) spec changed from (%s) to (%s) for cluster (%s), scaleResult=%+v", agentComponent, groupKey, currSpec, nextSpec, clusterName, scaleResult))

	_, err := controllerutil.CreateOrUpdate(ctx, pas.cli, &obj, func() error {
		if obj.Annotations == nil {
			obj.Annotations = make(map[string]string)
		}
		obj.Annotations[consts.AnnotationKeyO11yDesiredSpec] = nextSpec
		obj.Annotations[consts.AnnotationKeyO11ySpecUpdatedTimestamp] = strconv.Itoa(int(time.Now().Unix()))
		return nil
	})
	return err
}

func (pas *PodAutoScaler) updateSharedPoolAgentResources(ctx context.Context, clusterName, component string, scaleResult ScaleResult, specRange specRange) error {
	groupKey := scaleResult.GroupKey
	var obj agentv2.SharedPoolAgent
	if err := pas.cli.Get(ctx, client.ObjectKey{Namespace: groupKey.RefNamespace, Name: groupKey.RefName}, &obj); err != nil {
		return err
	}

	specs := pas.nextgenConfig.SharedPool.Specs
	agentComponent := consts.AgentSharedPool + "/" + component
	currSpec, nextSpec, changed, skipReason := decideSpecByScaleResult(&obj, scaleResult, specs, pas.autoscalerConfig.FrozenSeconds, specRange, component)
	if !changed {
		metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, skipReason})
		logger.Info(fmt.Sprintf("skip updating %s (%+v) spec (%s) for cluster (%s), reason=%s, scaleResult=%+v", agentComponent, groupKey, currSpec, clusterName, skipReason, scaleResult))
		return nil
	}

	metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, scaleResult.ActionResult.Action})
	logger.Info(fmt.Sprintf("%s (%+v) spec changed from (%s) to (%s) for cluster (%s), scaleResult=%+v", agentComponent, groupKey, currSpec, nextSpec, clusterName, scaleResult))

	_, err := controllerutil.CreateOrUpdate(ctx, pas.cli, &obj, func() error {
		if obj.Annotations == nil {
			obj.Annotations = make(map[string]string)
		}
		obj.Annotations[consts.AnnotationKeyO11yDesiredSpec] = combineSpec(&obj, nextSpec, component)
		obj.Annotations[consts.AnnotationKeyO11ySpecUpdatedTimestamp] = strconv.Itoa(int(time.Now().Unix()))
		return nil
	})
	return err
}

func (pas *PodAutoScaler) updateTidbGroupAgentResources(ctx context.Context, clusterName, component string, scaleResult ScaleResult, specRange specRange) error {
	groupKey := scaleResult.GroupKey
	var obj agentv2.TiDBGroupAgent
	if err := pas.cli.Get(ctx, client.ObjectKey{Namespace: groupKey.RefNamespace, Name: groupKey.RefName}, &obj); err != nil {
		return err
	}

	specs := pas.nextgenConfig.TiDBGroup.Specs
	agentComponent := consts.AgentTidbGroup + "/" + component
	currSpec, nextSpec, changed, skipReason := decideSpecByScaleResult(&obj, scaleResult, specs, pas.autoscalerConfig.FrozenSeconds, specRange, component)
	if !changed {
		metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, skipReason})
		logger.Info(fmt.Sprintf("skip updating %s (%+v) spec (%s) for cluster (%s), reason=%s, scaleResult=%+v", agentComponent, groupKey, currSpec, clusterName, skipReason, scaleResult))
		return nil
	}

	metrics.IncScaledResourceCounter([]string{agentComponent, clusterName, scaleResult.ActionResult.Action})
	logger.Info(fmt.Sprintf("%s (%+v) spec changed from (%s) to (%s) for cluster (%s), scaleResult=%+v", agentComponent, groupKey, currSpec, nextSpec, clusterName, scaleResult))

	_, err := controllerutil.CreateOrUpdate(ctx, pas.cli, &obj, func() error {
		if obj.Annotations == nil {
			obj.Annotations = make(map[string]string)
		}
		obj.Annotations[consts.AnnotationKeyO11yDesiredSpec] = nextSpec
		obj.Annotations[consts.AnnotationKeyO11ySpecUpdatedTimestamp] = strconv.Itoa(int(time.Now().Unix()))
		return nil
	})
	return err
}

// combineSpec combines the new spec with the existing desired spec in the annotations.
//
// The desired spec annotation uses the format "{vmagent_spec}/{vector_spec}" to support
// independent scaling of vmagent and vector components within the same agent.
// For example: "default/default", "large/small", or "/vector" (vmagent not set).
//
// It ensures that the spec for the other component remains unchanged when updating
// the spec for the specified component. If there is no existing desired spec,
// it initializes the annotation with the new spec for the specified component.
//
// Format specification:
//   - "{vmagent_spec}/{vector_spec}": both components specified
//   - "{vmagent_spec}/": only vmagent specified, vector empty
//   - "/{vector_spec}": only vector specified, vmagent empty
//   - "{single_spec}": legacy format, applies to both components
func combineSpec(obj metav1.Object, newSpec, component string) string {
	// 0: vmagent
	// 1: vector
	annotations := obj.GetAnnotations()
	desired := annotations[consts.AnnotationKeyO11yDesiredSpec]
	if desired == "" {
		if component == consts.ComponentVMAgent {
			return newSpec + "/"
		}
		return "/" + newSpec
	}

	// If the desired spec contains both vmagent and vector specs (separated by "/"),
	// we need to replace only the part corresponding to the specified component (vmagent or vector)
	// while keeping the other part unchanged.
	parts := strings.Split(desired, "/")
	switch len(parts) {
	case 2:
		if component == consts.ComponentVMAgent {
			return newSpec + "/" + parts[1]
		}
		return parts[0] + "/" + newSpec
	case 1:
		if component == consts.ComponentVMAgent {
			return newSpec + "/" + parts[0]
		}
		return parts[0] + "/" + newSpec
	}

	return newSpec
}

// decideSpecByScaleResult decides the new spec based on the scaling result and the annotations on the object.
// It checks the frozen window to avoid frequent changes.
// It uses the minSpec and maxSpec from annotations or falls back to the provided specRange.
// It returns the current spec, the new spec, whether the spec has changed, and the skip reason.
func decideSpecByScaleResult(
	obj metav1.Object,
	scaleResult ScaleResult,
	specs config.CustomerAgentSpecs,
	frozenSeconds int64,
	specRange specRange,
	component string,
) (string, string, bool, string) {
	annotations := obj.GetAnnotations()
	currentSpec := util.GetComponentDesiredSpec(annotations, defaultSpecName, component)

	lastUpdated, _ := strconv.Atoi(annotations[consts.AnnotationKeyO11ySpecUpdatedTimestamp])
	if time.Now().Unix()-int64(lastUpdated) < frozenSeconds {
		return currentSpec, "", false, skipReasonFrozen
	}

	// get minSpec and maxSpec from annotations or use defaults
	minSpec := util.GetComponentMinSpec(annotations, specRange.minSpec, component)
	maxSpec := util.GetComponentMaxSpec(annotations, specRange.maxSpec, component)

	result := selectSpec(specs, scaleResult.ActionResult.Action, currentSpec, minSpec, maxSpec)
	changed := result.newSpec != currentSpec

	return currentSpec, result.newSpec, changed, result.skipReason
}

// selectSpec selects the new spec based on the scaling action.
// Returns the selected spec and the reason if no change is made.
// For scaling up, it selects the next higher spec within the maxSpec limit.
// For scaling down, it selects the next lower spec within the minSpec limit.
func selectSpec(specs config.CustomerAgentSpecs, action string, currentSpec, minSpec, maxSpec string) specSelectionResult {
	var spec *config.CustomerAgentSpec

	switch action {
	case consts.ScaleActionUp:
		spec = specs.SelectUp(currentSpec, maxSpec)
		// Check if we're at maxSpec (SelectUp returns current spec when at max)
		if spec != nil && spec.Name == currentSpec {
			return specSelectionResult{
				newSpec:    currentSpec,
				skipReason: skipReasonMaxSpec,
			}
		}
	case consts.ScaleActionDown:
		spec = specs.SelectDown(currentSpec, minSpec)
		// Check if we're at minSpec (SelectDown returns current spec when at min)
		if spec != nil && spec.Name == currentSpec {
			return specSelectionResult{
				newSpec:    currentSpec,
				skipReason: skipReasonMinSpec,
			}
		}
	}

	// Spec not found or other error
	if spec == nil {
		return specSelectionResult{
			newSpec:    currentSpec,
			skipReason: skipReasonOther,
		}
	}

	// Successfully selected new spec
	return specSelectionResult{
		newSpec:    spec.Name,
		skipReason: "",
	}
}