package ratehandler

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

	"github.com/pingcap/dbaas/billing2/config"
	"github.com/pingcap/dbaas/billing2/utils/log"
	"github.com/pingcap/dbaas/billing2/utils/timeutils"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
	"golang.org/x/time/rate"
)

type Limiter interface {
	Allow(funcName string, args ...interface{}) bool
	SetLimit(funcName string, limit Limit)
}

type Limit struct {
	Value    int
	TimeSpan time.Duration
}

func PerMinute(val int) Limit {
	return Limit{
		Value:    val,
		TimeSpan: time.Minute,
	}
}

func buildArgsKey(args ...interface{}) string {
	strArgs := make([]string, len(args))
	for i, a := range args {
		strArgs[i] = fmt.Sprint(a)
	}
	return strings.Join(strArgs, "#")
}

type LocalLimit struct {
	mu           sync.Mutex
	limiterMap   map[string]map[string]*rate.Limiter
	funcLimitMap map[string]Limit
}

func (l *LocalLimit) Allow(funcName string, args ...interface{}) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	limit, ok := l.funcLimitMap[funcName]
	if !ok {
		return true
	}
	argsMap, ok := l.limiterMap[funcName]
	argsKey := buildArgsKey(args)
	argsLimiter, ok := argsMap[argsKey]
	if !ok {
		argsLimiter = rate.NewLimiter(rate.Every(limit.TimeSpan), limit.Value)
		l.limiterMap[funcName][argsKey] = argsLimiter
	}
	return argsLimiter.Allow()
}

func (l *LocalLimit) SetLimit(funcName string, limit Limit) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.funcLimitMap[funcName] = limit
	l.limiterMap[funcName] = make(map[string]*rate.Limiter)
}

func NewLocalLimiter() Limiter {
	return &LocalLimit{
		limiterMap:   make(map[string]map[string]*rate.Limiter),
		funcLimitMap: make(map[string]Limit),
	}
}

var globalLimiter = NewLocalLimiter()

func GetLimiter() Limiter {
	return globalLimiter
}

func SetLimiter(l Limiter) {
	globalLimiter = l
}

func Init() {
	r := redis.NewClient(&redis.Options{
		Addr: config.Get().GetString(config.RedisHost),
		DB:   config.Get().GetInt(config.RedisDB),
	})
	SetLimiter(NewRedisLimiter(r))
}

func NewRedisLimiter(r *redis.Client) Limiter {
	return &RedisLimiter{
		r: r,
	}
}

type RedisLimiter struct {
	r            *redis.Client
	funcLimitMap sync.Map
}

func (l *RedisLimiter) Allow(funcName string, args ...interface{}) bool {
	limiterInterface, ok := l.funcLimitMap.Load(funcName)
	if !ok {
		return true
	}
	limiter := limiterInterface.(Limit)
	timeKey := timeutils.GetNow().Truncate(limiter.TimeSpan).Unix()
	var redisKey string
	if len(args) == 0 {
		redisKey = fmt.Sprintf("counter-%d-func-%s-no-args", timeKey, funcName)
	} else {
		argsKey := buildArgsKey(args)
		redisKey = fmt.Sprintf("counter-%d-func-%s-args-%s", timeKey, funcName, argsKey)
	}
	ctx := context.Background()
	pipe := l.r.TxPipeline()
	incrResult := pipe.Incr(ctx, redisKey)
	pipe.Expire(ctx, redisKey, limiter.TimeSpan)
	_, err := pipe.Exec(ctx)
	if err != nil {
		log.Error("failed to inc key", zap.Error(err), zap.String("key", redisKey))
		return true
	}
	return int64(limiter.Value) >= incrResult.Val()
}

func (l *RedisLimiter) SetLimit(funcName string, limit Limit) {
	l.funcLimitMap.Store(funcName, limit)
}