package redis

import (
	"github.com/go-redis/redis"
	"strconv"
	"time"
)

type Z = redis.Z

type ZSet struct {
	key    string
	client *redis.Client
}

func NewZSet(client *redis.Client, key string) *ZSet {
	return &ZSet{
		key:    key,
		client: client,
	}
}

func (s *ZSet) Key() string {
	return s.key
}

// Expire 设置集合过期时间
func (s *ZSet) Expire(expiration time.Duration) (bool, error) {
	return s.client.Expire(s.key, expiration).Result()
}

// Add 添加一个或多个成员到有序集合
func (s *ZSet) Add(members ...Z) (int64, error) {
	return s.client.ZAdd(s.key, members...).Result()
}

// Count 返回有序集合在指定分数范围内的成员数量
func (s *ZSet) Count(min, max float64) (int64, error) {
	return s.client.ZCount(s.key, strconv.FormatFloat(min, 'f', -1, 64), strconv.FormatFloat(max, 'f', -1, 64)).Result()
}

// Score 获取成员的分数
func (s *ZSet) Score(member string) (float64, error) {
	return s.client.ZScore(s.key, member).Result()
}

// Rank 获取成员的排名(从小到大，0开始)
func (s *ZSet) Rank(member string) (int64, error) {
	return s.client.ZRank(s.key, member).Result()
}

// RevRank 获取成员的排名(从大到小，0开始)
func (s *ZSet) RevRank(member string) (int64, error) {
	return s.client.ZRevRank(s.key, member).Result()
}

// Range 返回指定排名范围的成员
func (s *ZSet) Range(start, stop int64) ([]string, error) {
	return s.client.ZRange(s.key, start, stop).Result()
}

// RangeWithScores 返回指定排名范围的成员和分数
func (s *ZSet) RangeWithScores(start, stop int64) ([]Z, error) {
	return s.client.ZRangeWithScores(s.key, start, stop).Result()
}

// RangeByScore 返回指定分数范围的成员
func (s *ZSet) RangeByScore(min, max float64) ([]string, error) {
	opt := redis.ZRangeBy{
		Min: strconv.FormatFloat(min, 'f', -1, 64),
		Max: strconv.FormatFloat(max, 'f', -1, 64),
	}
	return s.client.ZRangeByScore(s.key, opt).Result()
}

// RangeByScoreWithScores 返回指定分数范围的成员和分数
func (s *ZSet) RangeByScoreWithScores(min, max float64) ([]Z, error) {
	opt := redis.ZRangeBy{
		Min: strconv.FormatFloat(min, 'f', -1, 64),
		Max: strconv.FormatFloat(max, 'f', -1, 64),
	}
	return s.client.ZRangeByScoreWithScores(s.key, opt).Result()
}

// Remove 删除一个或多个成员
func (s *ZSet) Remove(members ...string) (int64, error) {
	return s.client.ZRem(s.key, members).Result()
}

// RemoveRangeByRank 删除指定排名范围的成员
func (s *ZSet) RemoveRangeByRank(start, stop int64) (int64, error) {
	return s.client.ZRemRangeByRank(s.key, start, stop).Result()
}

// RemoveRangeByScore 删除指定分数范围的成员
func (s *ZSet) RemoveRangeByScore(min, max float64) (int64, error) {
	return s.client.ZRemRangeByScore(s.key, strconv.FormatFloat(min, 'f', -1, 64), strconv.FormatFloat(max, 'f', -1, 64)).Result()
}

// Len 返回有序集合的成员数
func (s *ZSet) Len() (int64, error) {
	return s.client.ZCard(s.key).Result()
}

// IncScore 为成员的分数加上增量
func (s *ZSet) IncScore(member string, increment float64) (float64, error) {
	return s.client.ZIncrBy(s.key, increment, member).Result()
}

// Scan 扫描有序集合
func (s *ZSet) Scan(match string, count int64, callback func(member string, score float64) error) error {
	cursor := uint64(0)
	for {
		keys, newCursor, err := s.client.ZScan(s.key, cursor, match, count).Result()
		if err != nil {
			return err
		}

		for i := 0; i < len(keys); i += 2 {
			score, err := strconv.ParseFloat(keys[i+1], 64)
			if err != nil {
				return err
			}
			err = callback(keys[i], score)
			if err != nil {
				return err
			}
		}

		if newCursor == 0 {
			break
		}
		cursor = newCursor
	}
	return nil
}