package rcu

import (
	"context"
	"fmt"
	"math"
	"os"
	"strconv"
	"time"

	"github.com/pingcap/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/monitor"
	"github.com/tidbcloud/serverless-metering/pkg/repository"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"github.com/tidbcloud/serverless-metering/pkg/util"
	"go.uber.org/zap"
)

const (
	scaleMinute  = 5
	scaleSeconds = scaleMinute * 60
)

type RcuController struct {
	rcuStore               store.RcuStore
	repo                   repository.RcuMetaRepo
	localTime              util.LocalTime
	pushMetricAddr         string
	pushMetricInterval     int
	updateCapacityInterval int
}

func NewRcuController(rcuStore store.RcuStore, repo repository.RcuMetaRepo, localTime util.LocalTime) *RcuController {
	cfg := config.Global().RcuController
	return &RcuController{
		rcuStore:               rcuStore,
		repo:                   repo,
		localTime:              localTime,
		pushMetricAddr:         cfg.PushMetricAddr,
		pushMetricInterval:     cfg.PushMetricInterval,
		updateCapacityInterval: cfg.UpdateCapacityInterval,
	}
}

func (c *RcuController) Run() {
	log.Info("rcu controller started")
	c.updateProvisionedCapacity()
	ticker := time.NewTicker(time.Duration(c.updateCapacityInterval) * time.Second)
	defer ticker.Stop()
	metricPushTicker := time.NewTicker(time.Duration(c.pushMetricInterval) * time.Second)
	defer metricPushTicker.Stop()
	for {
		select {
		case <-metricPushTicker.C:
			c.pushCapacityMetrics()
		case <-ticker.C:
			c.updateProvisionedCapacity()
		}
	}
}

func (c *RcuController) pushCapacityMetrics() {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.pushMetricInterval)*time.Second)
	defer cancel()
	podName := os.Getenv("HOSTNAME")
	if podName == "" {
		podName = "default"
	}
	err := push.New(c.pushMetricAddr, "rcu-controller").Grouping("pod", podName).Collector(monitor.RcuCapacityGauge).PushContext(ctx)
	if err != nil {
		log.Error("failed to push rcu capacity metrics", zap.Error(err))
	}
}

func (c *RcuController) updateProvisionedCapacity() []*store.RcuCapacity {
	ctx := context.TODO()
	clusters, err := c.repo.ListEssentialPlanClusters(ctx)
	if err != nil {
		log.Error("failed to list essential plan clusters", zap.Error(err))
		monitor.RcuUpdateResult.WithLabelValues("ALL", "error").Inc()
		return nil
	}

	nowTs := c.localTime.Now().Unix()
	mapUsages, err := c.getRcuUsageMap(nowTs, clusters)
	if err != nil {
		log.Error("failed to get rcu usage map", zap.Error(err))
		monitor.RcuUpdateResult.WithLabelValues("ALL", "error").Inc()
		return nil
	}

	mapCapacities, err := c.getRcuCapacityMap()
	if err != nil {
		log.Error("failed to get rcu capacity map", zap.Error(err))
		monitor.RcuUpdateResult.WithLabelValues("ALL", "error").Inc()
		return nil
	}

	result := make([]*store.RcuCapacity, 0)
	for _, cluster := range clusters {
		capacity, err := c.updateProvisionedCapacityForCluster(cluster, mapUsages, mapCapacities)
		if err != nil {
			log.Error("failed to update provisioned capacity for cluster", zap.String("clusterId", cluster.PrimaryID), zap.Error(err))
			monitor.RcuUpdateResult.WithLabelValues(fmt.Sprintf("%d", cluster.KeyspaceID), "error").Inc()
			continue
		} else {
			monitor.RcuUpdateResult.WithLabelValues(fmt.Sprintf("%d", cluster.KeyspaceID), "success").Inc()
		}
		result = append(result, capacity)
	}

	c.deleteExtraRcuCapacity(clusters, mapCapacities)

	return result
}

func (c *RcuController) updateProvisionedCapacityForCluster(
	cluster *repository.ClusterBranch,
	mapUsages map[uint32]*store.RcuUsage,
	mapCapacities map[uint32]*store.RcuCapacity,
) (*store.RcuCapacity, error) {
	logger := log.With(zap.String("clusterId", cluster.PrimaryID))
	if cluster.Type != repository.CBTypeCluster {
		return nil, fmt.Errorf("cluster type is not cluster, type: %d", cluster.Type)
	}
	if !cluster.IsEssentialPlan() {
		return nil, fmt.Errorf("cluster is not essential plan, cluster id: %s", cluster.PrimaryID)
	}
	monitor.RcuMinGauge.WithLabelValues(fmt.Sprintf("%d", cluster.KeyspaceID), cluster.KeyspaceToken, cluster.PrimaryID).Set(float64(cluster.MinRcu))
	monitor.RcuMaxGauge.WithLabelValues(fmt.Sprintf("%d", cluster.KeyspaceID), cluster.KeyspaceToken, cluster.PrimaryID).Set(float64(cluster.MaxRcu))

	usage, ok := mapUsages[cluster.KeyspaceID]
	// The cluster hasn't been used in past 15 minutes, use min rcu as current capacity
	if !ok {
		return c.setCapacity(cluster, cluster.MinRcu)
	} else if len(usage.Values) == 0 {
		return c.setCapacity(cluster, cluster.MinRcu)
	}

	capacity, ok := mapCapacities[cluster.KeyspaceID]
	// The cluster capacity hasn't been stored, use min rcu as current capacity
	if !ok {
		return c.setCapacity(cluster, cluster.MinRcu)
	}

	if len(usage.Values) < scaleMinute {
		logger.Info("usage values less than scale minute", zap.Int("scaleMinute", scaleMinute), zap.Int64s("usageValues", usage.Values))
		return c.setCapacity(cluster, capacity.ProvisionedCapacity)
	}

	targetCapacity := int64(math.Ceil(util.PercentileInt64(usage.Values, 75) * 1.1))
	if targetCapacity < cluster.MinRcu {
		targetCapacity = cluster.MinRcu
	} else if targetCapacity > cluster.MaxRcu {
		targetCapacity = cluster.MaxRcu
	}
	logger.Info("updateProvisionedCapacity", zap.Int64("targetCapacity", targetCapacity), zap.Int64s("usageValues", usage.Values))
	return c.setCapacity(cluster, targetCapacity)
}

func (c *RcuController) setCapacity(cluster *repository.ClusterBranch, provisioned int64) (*store.RcuCapacity, error) {
	capacity := &store.RcuCapacity{
		KeyspaceID:          cluster.KeyspaceID,
		ClusterID:           cluster.ClusterID,
		ProvisionedCapacity: provisioned,
		MinRcu:              cluster.MinRcu,
		MaxRcu:              cluster.MaxRcu,
	}
	err := c.rcuStore.WriteRcuCapacity(capacity)
	if err != nil {
		return nil, err
	}
	monitor.RcuCapacityGauge.WithLabelValues(fmt.Sprintf("%d", cluster.KeyspaceID), cluster.KeyspaceToken, cluster.PrimaryID).Set(float64(capacity.ProvisionedCapacity))
	log.Info("set capacity",
		zap.String("clusterId", cluster.PrimaryID),
		zap.Uint32("keyspace_id", cluster.KeyspaceID),
		zap.Int64("provisionedCapacity", capacity.ProvisionedCapacity),
		zap.Int64("minRcu", capacity.MinRcu),
		zap.Int64("maxRcu", capacity.MaxRcu),
	)
	return capacity, nil
}

func (c *RcuController) getRcuCapacityMap() (map[uint32]*store.RcuCapacity, error) {
	capacities, err := c.rcuStore.ListRcuCapacities()
	if err != nil {
		return nil, err
	}
	mapCapacities := make(map[uint32]*store.RcuCapacity)
	for _, capacity := range capacities {
		mapCapacities[capacity.KeyspaceID] = capacity
	}
	return mapCapacities, nil
}

func (c *RcuController) deleteExtraRcuCapacity(clusters []*repository.ClusterBranch, capacityMap map[uint32]*store.RcuCapacity) {
	for _, cluster := range clusters {
		delete(capacityMap, cluster.KeyspaceID)
	}
	for keyspaceID := range capacityMap {
		err := c.rcuStore.DeleteRcuCapacity(keyspaceID)
		if err != nil {
			log.Error("failed to delete rcu capacity", zap.Uint32("keyspaceId", keyspaceID), zap.Error(err))
			monitor.RcuUpdateResult.WithLabelValues(fmt.Sprintf("%d", keyspaceID), "error").Inc()
			continue
		}

		monitor.RcuCapacityGauge.DeletePartialMatch(prometheus.Labels{"keyspace_id": fmt.Sprintf("%d", keyspaceID)})
		monitor.RcuMinGauge.DeletePartialMatch(prometheus.Labels{"keyspace_id": fmt.Sprintf("%d", keyspaceID)})
		monitor.RcuMaxGauge.DeletePartialMatch(prometheus.Labels{"keyspace_id": fmt.Sprintf("%d", keyspaceID)})
		monitor.RcuUpdateResult.WithLabelValues(fmt.Sprintf("%d", keyspaceID), "success").Inc()
		log.Info("deleted rcu capacity", zap.Uint32("keyspaceId", keyspaceID))
	}
}

func (c *RcuController) getRcuUsageMap(nowTs int64, clusters []*repository.ClusterBranch) (map[uint32]*store.RcuUsage, error) {
	// latest on the minute
	now := util.LocateLatestOnTheMinute(nowTs)
	// 5 minutes ago
	prev := util.LocateLatestOnTheMinute(nowTs - scaleSeconds)
	mapUsages := make(map[uint32]*store.RcuUsage)

	var keyspaceIDsV1 []uint32
	var keyspaceIDsV2 []uint32
	for _, cluster := range clusters {
		if cluster.UseRUV2() {
			keyspaceIDsV2 = append(keyspaceIDsV2, cluster.KeyspaceID)
			continue
		}
		keyspaceIDsV1 = append(keyspaceIDsV1, cluster.KeyspaceID)
	}

	appendUsages := func(usages []*store.RcuUsage) {
		for _, usage := range usages {
			keyspaceID, err := strconv.ParseInt(usage.KeyspaceID, 10, 64)
			if err != nil {
				log.Error("failed to parse keyspace id", zap.String("keyspaceId", usage.KeyspaceID), zap.Error(err))
				continue
			}
			mapUsages[uint32(keyspaceID)] = usage
		}
	}

	if len(keyspaceIDsV1) > 0 {
		usagesV1, err := c.rcuStore.CalRcuUsages(keyspaceIDsV1, prev, now, false)
		if err != nil {
			return nil, err
		}
		appendUsages(usagesV1)
	}

	if len(keyspaceIDsV2) > 0 {
		usagesV2, err := c.rcuStore.CalRcuUsages(keyspaceIDsV2, prev, now, true)
		if err != nil {
			return nil, err
		}
		appendUsages(usagesV2)
	}

	return mapUsages, nil
}