package service

import (
	"context"
	"fmt"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

var acquireLua = `
local current = tonumber(redis.call("GET", KEYS[1]) or "0")
local max = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
if current >= max then
  return 0
end
current = redis.call("INCR", KEYS[1])
if current == 1 then
  redis.call("EXPIRE", KEYS[1], ttl)
end
if current > max then
  redis.call("DECR", KEYS[1])
  return 0
end
return 1
`

var releaseLua = `
local current = tonumber(redis.call("DECR", KEYS[1]) or "0")
if current <= 0 then
  redis.call("DEL", KEYS[1])
end
return 1
`

func acquireSemaphore(ctx context.Context, key string, max int, ttlSeconds int64) (bool, error) {
	if max <= 0 {
		// 不限制
		return true, nil
	}
	if ttlSeconds <= 0 {
		ttlSeconds = 3600
	}
	r, err := g.Redis().Do(ctx, "EVAL", acquireLua, 1, key, max, ttlSeconds)
	if err != nil {
		return false, fmt.Errorf("获取并发令牌失败: %w", err)
	}
	return gconv.Int(r) == 1, nil
}

func releaseSemaphore(ctx context.Context, key string) error {
	_, err := g.Redis().Do(ctx, "EVAL", releaseLua, 1, key)
	return err
}