package goutil

import (
	"fmt"
	"math/rand"
	"time"

	"github.com/go-redis/redis"
)

type WindowsLimitOption struct {
	Key     string        //redis可以前缀，会拼接上succ和fail
	Count   int64         //限制数量
	Size    time.Duration //限制时间区间，默认一分钟
	succKey string
	failKey string
}
type WindowsLimit struct {
	client *redis.Client
	option *WindowsLimitOption
}

func NewWindowsLimit(client *redis.Client, option *WindowsLimitOption) *WindowsLimit {
	option.Default()
	return &WindowsLimit{
		client: client,
		option: option,
	}
}

func (o *WindowsLimitOption) Default() {
	if o.Key == "" {
		o.Key = fmt.Sprintf("windows:limit:%s:%d", time.Now().Format("20060102"), time.Now().UnixMilli())
	}
	if o.Count <= 0 {
		o.Count = 1000
	}
	if o.Size <= 0 {
		o.Size = time.Second * 60
	}
}
func (w *WindowsLimit) tidy() {
	now := time.Now().UnixMilli()
	start := now - w.option.Size.Milliseconds()

	pipe := w.client.Pipeline()
	pipe.ZRemRangeByScore(w.option.failKey, "0", fmt.Sprintf("%d", start))
	pipe.ZRemRangeByScore(w.option.succKey, "0", fmt.Sprintf("%d", start))
	pipe.Exec()
}

func (w *WindowsLimit) Count() (int64, int64, error) {
	now := time.Now().UnixMilli()
	start := now - w.option.Size.Milliseconds()

	pipe := w.client.Pipeline()
	succCmd := pipe.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
	failCmd := pipe.ZCount(w.option.failKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
	_, err := pipe.Exec()
	if err != nil {
		return 0, 0, err
	}
	return succCmd.Val(), failCmd.Val(), nil
}

func (w *WindowsLimit) Allow() bool {
	now := time.Now().UnixMilli()
	start := now - w.option.Size.Milliseconds()
	pipe := w.client.Pipeline()
	succCmd := w.client.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
	_, err := pipe.Exec()
	if err != nil {
		return true
	}

	if succCmd.Val() <= w.option.Count {
		return true
	}
	return false
}

func (w *WindowsLimit) Add(succ bool) error {
	now := time.Now().UnixMilli()
	w.tidy()
	pipe := w.client.Pipeline()
	if succ {
		pipe.ZAdd(w.option.succKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
	} else {
		pipe.ZAdd(w.option.failKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
	}

	pipe.Expire(w.option.failKey, w.option.Size)
	pipe.Expire(w.option.succKey, w.option.Size)
	_, err := pipe.Exec()
	return err
}