package redis

import (
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

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

//go:embed limit.lua
var lua string

const (
	SUCCESS = "success"
	FAILED  = "failed"
)

type Bucket struct {
	Name        string          `json:"na"`
	Capacity    int64           `json:"ca"`
	Rate        int64           `json:"ra"`
	Size        int64           `json:"si"`
	GenerateAt  int64           `json:"ge_at"`
	CreatedAt   int64           `json:"cr_at"`
	UpdatedAt   int64           `json:"up_at"`
	PassedCount int64           `json:"pa_co"`
	RejectCount int64           `json:"re_co"`
	PassedHours map[int64]int64 `json:"pa_ho"`
	RejectHours map[int64]int64 `json:"re_ho"`
}

type BucketList []*Bucket

func (a BucketList) Len() int { return len(a) }
func (a BucketList) Less(i, j int) bool {
	return a[i].PassedCount > a[j].PassedCount
}
func (a BucketList) Swap(i, j int) { a[i], a[j] = a[j], a[i] }

type TokenBucket struct {
	keyGen   func() string
	capacity int64
	rate     int64
	script   *redis.Script
	client   *redis.Client
	expire   time.Duration
	lock     sync.RWMutex
	curKey   string
}

func NewTokenBucket(client *redis.Client, keyGen func() string, capacity, rate int64, expire time.Duration) *TokenBucket {
	return &TokenBucket{
		expire:   expire,
		capacity: capacity,
		rate:     rate,
		client:   client,
		script:   redis.NewScript(lua),
		keyGen:   keyGen,
	}
}

func (s *TokenBucket) TakeTokens(name string, count int64) error {
	key := s.keyGen()
	now := time.Now().Local()
	cmd := []string{
		key,  //key
		name, //name
		fmt.Sprintf("%d", s.capacity),
		fmt.Sprintf("%d", s.rate),
		fmt.Sprintf("%d", count),
		fmt.Sprintf("%d", now.UnixMilli()),
		fmt.Sprintf("%d", now.Hour()),
	}
	ret, err := s.script.Run(s.client, cmd).String()
	if err != nil {
		return err
	}
	if ret != SUCCESS {
		return errors.New("sorry, I can't work anymore~")
	}
	var same = false
	if s.lock.TryRLock() {
		same = key == s.curKey
		s.lock.RUnlock()
	} else {
		return nil
	}
	if !same {
		if s.lock.TryLock() {
			if ok, err := s.client.Expire(key, s.expire).Result(); err == nil && ok {
				s.curKey = key
			}
			s.lock.Unlock()
		}
	}
	return nil
}

func (s *TokenBucket) Traverse(callback func(key string, item *Bucket)) error {
	return NewHash(s.client, s.keyGen()).Scan("", 1000, func(name, value string) {
		bucket := Bucket{}
		if err := json.Unmarshal([]byte(value), &bucket); err == nil {
			callback(name, &bucket)
		}
	})
}