// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package nodes

import (
	"context"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
	"github.com/TeaOSLab/EdgeNode/internal/events"
	"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
	"sync"
)

func init() {
	if !teaconst.IsMain {
		return
	}

	events.On(events.EventLoaded, func() {
		var cfg = nodeConfig()
		if cfg == nil {
			return
		}
		sharedPlanBandwidthLimiter.UpdatePlans(cfg.FindAllPlans())
	})
}

var sharedPlanBandwidthLimiter = NewPlanBandwidthLimiter()

// PlanBandwidthLimiter 套餐带宽限制
type PlanBandwidthLimiter struct {
	m map[int64]*ratelimit.Bandwidth // id => *Bandwidth

	mu *sync.RWMutex
}

// NewPlanBandwidthLimiter 获取新的Limiter
func NewPlanBandwidthLimiter() *PlanBandwidthLimiter {
	return &PlanBandwidthLimiter{
		mu: &sync.RWMutex{},
		m:  make(map[int64]*ratelimit.Bandwidth),
	}
}

// UpdatePlans 更新套餐信息
func (this *PlanBandwidthLimiter) UpdatePlans(planMap map[int64]*serverconfigs.PlanConfig) {
	var bandwidthMap = map[int64]*ratelimit.Bandwidth{}
	for _, plan := range planMap {
		if plan.BandwidthLimitPerNode == nil || plan.BandwidthLimitPerNode.Count <= 0 {
			continue
		}
		var bits = plan.BandwidthLimitPerNode.Bits()
		if bits <= 0 {
			continue
		}

		var bytes = bits / 8
		if bytes == 0 {
			continue
		}
		bandwidthMap[plan.Id] = ratelimit.NewBandwidth(bytes)
	}

	this.mu.Lock()
	this.m = bandwidthMap
	this.mu.Unlock()
}

// Ack 获取下一次发送数据的机会
func (this *PlanBandwidthLimiter) Ack(ctx context.Context, planId int64, newBytes int) {
	if planId <= 0 {
		return
	}

	this.mu.RLock()
	defer this.mu.RUnlock()
	bandwidth, ok := this.m[planId]
	if !ok {
		return
	}
	bandwidth.Ack(ctx, newBytes)
}